YARN-8103. Add CLI interface to query node attributes. Contributed by Bibin A Chundatt.
Project: http://git-wip-us.apache.org/repos/asf/hadoop/repo Commit: http://git-wip-us.apache.org/repos/asf/hadoop/commit/acd7729a Tree: http://git-wip-us.apache.org/repos/asf/hadoop/tree/acd7729a Diff: http://git-wip-us.apache.org/repos/asf/hadoop/diff/acd7729a Branch: refs/heads/YARN-3409 Commit: acd7729af570ae2c59a547709d510bf20b0b1246 Parents: 649e014 Author: Naganarasimha <naganarasimha...@apache.org> Authored: Thu Jun 28 08:13:09 2018 +0800 Committer: Sunil G <sun...@apache.org> Committed: Fri Jul 27 17:01:44 2018 +0530 ---------------------------------------------------------------------- .../hadoop/yarn/sls/nodemanager/NodeInfo.java | 12 +- .../yarn/sls/scheduler/RMNodeWrapper.java | 8 +- hadoop-yarn-project/hadoop-yarn/bin/yarn | 4 +- .../hadoop/yarn/api/records/NodeReport.java | 13 + .../src/main/proto/yarn_protos.proto | 1 + .../hadoop/yarn/client/cli/ClusterCLI.java | 17 + .../yarn/client/cli/NodeAttributesCLI.java | 893 +++++++++++++------ .../apache/hadoop/yarn/client/cli/NodeCLI.java | 13 +- .../hadoop/yarn/client/cli/TestClusterCLI.java | 32 +- .../yarn/client/cli/TestNodeAttributesCLI.java | 331 +++++-- .../hadoop/yarn/client/cli/TestYarnCLI.java | 31 +- .../impl/pb/NodeAttributeInfoPBImpl.java | 10 +- .../records/impl/pb/NodeAttributePBImpl.java | 12 +- .../api/records/impl/pb/NodeReportPBImpl.java | 44 +- .../hadoop/yarn/server/utils/BuilderUtils.java | 6 +- .../server/resourcemanager/AdminService.java | 5 +- .../server/resourcemanager/ClientRMService.java | 5 +- .../resourcemanager/ResourceTrackerService.java | 4 - .../nodelabels/NodeAttributesManagerImpl.java | 2 +- .../server/resourcemanager/rmnode/RMNode.java | 11 +- .../resourcemanager/rmnode/RMNodeImpl.java | 18 +- .../resourcemanager/webapp/dao/NodeInfo.java | 13 +- .../yarn/server/resourcemanager/MockNodes.java | 4 +- .../resourcemanager/TestRMAdminService.java | 14 +- .../clientrm/FederationClientInterceptor.java | 6 +- 25 files changed, 1053 insertions(+), 456 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java ---------------------------------------------------------------------- diff --git a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java index 65b8da0..2eee351 100644 --- a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java +++ b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/nodemanager/NodeInfo.java @@ -19,6 +19,7 @@ package org.apache.hadoop.yarn.sls.nodemanager; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Set; @@ -220,16 +221,9 @@ public class NodeInfo { return null; } - - @Override - public void setNodeAttributes(String prefix, - Set<NodeAttribute> nodeAttributes) { - - } - @Override - public Map<String, Set<NodeAttribute>> getAllNodeAttributes() { - return null; + public Set<NodeAttribute> getAllNodeAttributes() { + return Collections.emptySet(); } @Override http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java ---------------------------------------------------------------------- diff --git a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java index bf61f54..248b634 100644 --- a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java +++ b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/scheduler/RMNodeWrapper.java @@ -209,13 +209,7 @@ public class RMNodeWrapper implements RMNode { } @Override - public void setNodeAttributes(String prefix, - Set<NodeAttribute> nodeAttributes) { - node.setNodeAttributes(prefix, nodeAttributes); - } - - @Override - public Map<String, Set<NodeAttribute>> getAllNodeAttributes() { + public Set<NodeAttribute> getAllNodeAttributes() { return node.getAllNodeAttributes(); } http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/bin/yarn ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/bin/yarn b/hadoop-yarn-project/hadoop-yarn/bin/yarn index 7cd838f..8290fcd 100755 --- a/hadoop-yarn-project/hadoop-yarn/bin/yarn +++ b/hadoop-yarn-project/hadoop-yarn/bin/yarn @@ -55,7 +55,7 @@ function hadoop_usage hadoop_add_subcommand "timelinereader" client "run the timeline reader server" hadoop_add_subcommand "timelineserver" daemon "run the timeline server" hadoop_add_subcommand "top" client "view cluster information" - hadoop_add_subcommand "node-attributes" "map node to attibutes" + hadoop_add_subcommand "nodeattributes" client "node attributes cli client" hadoop_add_subcommand "version" client "print the version" hadoop_generate_usage "${HADOOP_SHELL_EXECNAME}" true } @@ -187,7 +187,7 @@ ${HADOOP_COMMON_HOME}/${HADOOP_COMMON_LIB_JARS_DIR}" hadoop_add_classpath "$HADOOP_YARN_HOME/$YARN_DIR/timelineservice/lib/*" HADOOP_CLASSNAME='org.apache.hadoop.yarn.server.timelineservice.reader.TimelineReaderServer' ;; - node-attributes) + nodeattributes) HADOOP_SUBCMD_SUPPORTDAEMONIZATION="false" HADOOP_CLASSNAME='org.apache.hadoop.yarn.client.cli.NodeAttributesCLI' ;; http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeReport.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeReport.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeReport.java index 3a80641..625ad23 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeReport.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeReport.java @@ -258,4 +258,17 @@ public abstract class NodeReport { * Set the node update type (null indicates absent node update type). * */ public void setNodeUpdateType(NodeUpdateType nodeUpdateType) {} + + /** + * Set the node attributes of node. + * + * @param nodeAttributes set of node attributes. + */ + public abstract void setNodeAttributes(Set<NodeAttribute> nodeAttributes); + + /** + * Get node attributes of node. + * @return the set of node attributes. + */ + public abstract Set<NodeAttribute> getNodeAttributes(); } http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto index aca9471..10b36c7 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto @@ -355,6 +355,7 @@ message NodeReportProto { optional ResourceUtilizationProto node_utilization = 12; optional uint32 decommissioning_timeout = 13; optional NodeUpdateTypeProto node_update_type = 14; + repeated NodeAttributeProto node_attributes = 15; } message NodeIdToLabelsProto { http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java index a29b0db..4d93949 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java @@ -36,6 +36,7 @@ import org.apache.commons.lang3.StringUtils; import org.apache.hadoop.classification.InterfaceAudience.Private; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.util.ToolRunner; +import org.apache.hadoop.yarn.api.records.NodeAttributeInfo; import org.apache.hadoop.yarn.api.records.NodeLabel; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; @@ -52,6 +53,7 @@ public class ClusterCLI extends YarnCLI { public static final String LIST_LABELS_CMD = "list-node-labels"; public static final String DIRECTLY_ACCESS_NODE_LABEL_STORE = "directly-access-node-label-store"; + public static final String LIST_CLUSTER_ATTRIBUTES="list-node-attributes"; public static final String CMD = "cluster"; private boolean accessLocal = false; static CommonNodeLabelsManager localNodeLabelsManager = null; @@ -71,6 +73,8 @@ public class ClusterCLI extends YarnCLI { opts.addOption("lnl", LIST_LABELS_CMD, false, "List cluster node-label collection"); + opts.addOption("lna", LIST_CLUSTER_ATTRIBUTES, false, + "List cluster node-attribute collection"); opts.addOption("h", HELP_CMD, false, "Displays help for all commands."); opts.addOption("dnl", DIRECTLY_ACCESS_NODE_LABEL_STORE, false, "This is DEPRECATED, will be removed in future releases. Directly access node label store, " @@ -102,6 +106,8 @@ public class ClusterCLI extends YarnCLI { if (parsedCli.hasOption(LIST_LABELS_CMD)) { printClusterNodeLabels(); + } else if(parsedCli.hasOption(LIST_CLUSTER_ATTRIBUTES)){ + printClusterNodeAttributes(); } else if (parsedCli.hasOption(HELP_CMD)) { printUsage(opts); return 0; @@ -112,6 +118,17 @@ public class ClusterCLI extends YarnCLI { return 0; } + private void printClusterNodeAttributes() throws IOException, YarnException { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter pw = new PrintWriter( + new OutputStreamWriter(baos, Charset.forName("UTF-8"))); + for (NodeAttributeInfo attribute : client.getClusterAttributes()) { + pw.println(attribute.toString()); + } + pw.close(); + sysout.println(baos.toString("UTF-8")); + } + void printClusterNodeLabels() throws YarnException, IOException { List<NodeLabel> nodeLabels = null; if (accessLocal) { http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java index df5a57d..13d5e24 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java @@ -18,29 +18,30 @@ package org.apache.hadoop.yarn.client.cli; -import java.io.IOException; -import java.io.PrintStream; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; - +import com.google.common.base.Preconditions; +import com.google.common.collect.Lists; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.GnuParser; import org.apache.commons.cli.MissingArgumentException; +import org.apache.commons.cli.Option; +import org.apache.commons.cli.OptionGroup; import org.apache.commons.cli.Options; -import org.apache.commons.cli.ParseException; +import org.apache.commons.cli.UnrecognizedOptionException; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.conf.Configured; import org.apache.hadoop.fs.CommonConfigurationKeys; -import org.apache.hadoop.ha.HAAdmin.UsageInfo; -import org.apache.hadoop.ipc.RemoteException; import org.apache.hadoop.util.Tool; import org.apache.hadoop.util.ToolRunner; +import org.apache.hadoop.yarn.api.ApplicationClientProtocol; +import org.apache.hadoop.yarn.api.protocolrecords.GetAttributesToNodesRequest; +import org.apache.hadoop.yarn.api.protocolrecords.GetAttributesToNodesResponse; +import org.apache.hadoop.yarn.api.protocolrecords.GetClusterNodeAttributesRequest; +import org.apache.hadoop.yarn.api.protocolrecords.GetClusterNodeAttributesResponse; +import org.apache.hadoop.yarn.api.protocolrecords.GetNodesToAttributesRequest; +import org.apache.hadoop.yarn.api.protocolrecords.GetNodesToAttributesResponse; import org.apache.hadoop.yarn.api.records.NodeAttribute; +import org.apache.hadoop.yarn.api.records.NodeAttributeInfo; +import org.apache.hadoop.yarn.api.records.NodeAttributeKey; import org.apache.hadoop.yarn.api.records.NodeAttributeType; import org.apache.hadoop.yarn.client.ClientRMProxy; import org.apache.hadoop.yarn.conf.YarnConfiguration; @@ -50,13 +51,24 @@ import org.apache.hadoop.yarn.server.api.protocolrecords.AttributeMappingOperati import org.apache.hadoop.yarn.server.api.protocolrecords.NodeToAttributes; import org.apache.hadoop.yarn.server.api.protocolrecords.NodesToAttributesMappingRequest; -import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Lists; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStreamWriter; +import java.io.PrintStream; +import java.io.PrintWriter; +import java.io.UnsupportedEncodingException; +import java.nio.charset.Charset; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Set; /** * CLI to map attributes to Nodes. - * */ public class NodeAttributesCLI extends Configured implements Tool { @@ -64,351 +76,640 @@ public class NodeAttributesCLI extends Configured implements Tool { "Invalid Node to attribute mapping : "; protected static final String USAGE_YARN_NODE_ATTRIBUTES = - "Usage: yarn node-attributes "; + "Usage: yarn nodeattributes "; + + protected static final String MISSING_ARGUMENT = + "Missing argument for command"; protected static final String NO_MAPPING_ERR_MSG = "No node-to-attributes mappings are specified"; - protected final static Map<String, UsageInfo> NODE_ATTRIB_USAGE = - ImmutableMap.<String, UsageInfo>builder() - .put("-replace", - new UsageInfo( - "<\"node1:attribute[(type)][=value],attribute1[=value]," - + "attribute2 node2:attribute2[=value],attribute3\">", - " Replace the node to attributes mapping information at the" - + " ResourceManager with the new mapping. Currently" - + " supported attribute type. And string is the default" - + " type too. Attribute value if not specified for string" - + " type value will be considered as empty string." - + " Replaced node-attributes should not violate the" - + " existing attribute to attribute type mapping.")) - .put("-add", - new UsageInfo( - "<\"node1:attribute[(type)][=value],attribute1[=value]," - + "attribute2 node2:attribute2[=value],attribute3\">", - " Adds or updates the node to attributes mapping information" - + " at the ResourceManager. Currently supported attribute" - + " type is string. And string is the default type too." - + " Attribute value if not specified for string type" - + " value will be considered as empty string. Added or" - + " updated node-attributes should not violate the" - + " existing attribute to attribute type mapping.")) - .put("-remove", - new UsageInfo("<\"node1:attribute,attribute1 node2:attribute2\">", - " Removes the specified node to attributes mapping" - + " information at the ResourceManager")) - .put("-failOnUnknownNodes", - new UsageInfo("", - "Can be used optionally along with other options. When its" - + " set, it will fail if specified nodes are unknown.")) - .build(); - - /** Output stream for errors, for use in tests. */ + private static final String DEFAULT_SEPARATOR = System.lineSeparator(); + public static final String INVALID_COMMAND_USAGE = "Invalid Command Usage : "; + /** + * Output stream for errors, for use in tests. + */ private PrintStream errOut = System.err; public NodeAttributesCLI() { super(); } - public NodeAttributesCLI(Configuration conf) { - super(conf); - } - protected void setErrOut(PrintStream errOut) { this.errOut = errOut; } - private void printHelpMsg(String cmd) { - StringBuilder builder = new StringBuilder(); - UsageInfo usageInfo = null; - if (cmd != null && !(cmd.trim().isEmpty())) { - usageInfo = NODE_ATTRIB_USAGE.get(cmd); - } - if (usageInfo != null) { - if (usageInfo.args == null) { - builder.append(" " + cmd + ":\n" + usageInfo.help); - } else { - String space = (usageInfo.args == "") ? "" : " "; - builder.append( - " " + cmd + space + usageInfo.args + " :\n" + usageInfo.help); - } - } else { - // help for all commands - builder.append("Usage: yarn node-attributes\n"); - for (Map.Entry<String, UsageInfo> cmdEntry : NODE_ATTRIB_USAGE - .entrySet()) { - usageInfo = cmdEntry.getValue(); - builder.append(" " + cmdEntry.getKey() + " " + usageInfo.args - + " :\n " + usageInfo.help + "\n"); - } - builder.append(" -help" + " [cmd]\n"); - } - errOut.println(builder); + protected AdminCommandHandler getAdminCommandHandler() { + return new AdminCommandHandler(); } - private static void buildIndividualUsageMsg(String cmd, - StringBuilder builder) { - UsageInfo usageInfo = NODE_ATTRIB_USAGE.get(cmd); - if (usageInfo == null) { - return; - } - if (usageInfo.args == null) { - builder.append(USAGE_YARN_NODE_ATTRIBUTES + cmd + "\n"); - } else { - String space = (usageInfo.args == "") ? "" : " "; - builder.append( - USAGE_YARN_NODE_ATTRIBUTES + cmd + space + usageInfo.args + "\n"); - } + protected ClientCommandHandler getClientCommandHandler() { + return new ClientCommandHandler(); } - private static void buildUsageMsgForAllCmds(StringBuilder builder) { - builder.append("Usage: yarn node-attributes\n"); - for (Map.Entry<String, UsageInfo> cmdEntry : NODE_ATTRIB_USAGE.entrySet()) { - UsageInfo usageInfo = cmdEntry.getValue(); - builder.append(" " + cmdEntry.getKey() + " " + usageInfo.args + "\n"); + void printUsage(String cmd, boolean desc, CommandHandler... handlers) + throws UnsupportedEncodingException { + StringBuilder usageBuilder = new StringBuilder(); + usageBuilder.append(USAGE_YARN_NODE_ATTRIBUTES); + boolean satisfied = false; + for (CommandHandler cmdHandlers : handlers) { + satisfied |= cmdHandlers.getHelp(cmd, usageBuilder, desc); + } + if (!satisfied) { + printUsage(desc, handlers); + } else { + print(usageBuilder); } - builder.append(" -help" + " [cmd]\n"); } - /** - * Displays format of commands. - * - * @param cmd The command that is being executed. - */ - private void printUsage(String cmd) { + private void printUsage(boolean desc, CommandHandler... handlers) + throws UnsupportedEncodingException { StringBuilder usageBuilder = new StringBuilder(); - if (NODE_ATTRIB_USAGE.containsKey(cmd)) { - buildIndividualUsageMsg(cmd, usageBuilder); - } else { - buildUsageMsgForAllCmds(usageBuilder); + usageBuilder.append(USAGE_YARN_NODE_ATTRIBUTES); + for (CommandHandler cmdHandlers : handlers) { + cmdHandlers.getHelp(usageBuilder, desc); } - errOut.println(usageBuilder); - } - private void printUsage() { - printUsage(""); + // append help with usage + usageBuilder.append(DEFAULT_SEPARATOR) + .append(" -help [cmd] List help of commands"); + print(usageBuilder); } - protected ResourceManagerAdministrationProtocol createAdminProtocol() - throws IOException { - // Get the current configuration - final YarnConfiguration conf = new YarnConfiguration(getConf()); - return ClientRMProxy.createRMProxy(conf, - ResourceManagerAdministrationProtocol.class); + private void print(StringBuilder usageBuilder) + throws UnsupportedEncodingException { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter pw = + new PrintWriter(new OutputStreamWriter(baos, Charset.forName("UTF-8"))); + pw.write(usageBuilder.toString()); + pw.close(); + errOut.println(baos.toString("UTF-8")); } - @Override - public void setConf(Configuration conf) { - if (conf != null) { - conf = addSecurityConfiguration(conf); + private Options buildOptions(CommandHandler... handlers) { + Options opts = new Options(); + for (CommandHandler handler : handlers) { + Options handlerOpts = handler.getOptions(); + handlerOpts.getOptions().iterator() + .forEachRemaining(option -> opts.addOption((Option) option)); } - super.setConf(conf); - } - - /** - * Add the requisite security principal settings to the given Configuration, - * returning a copy. - * - * @param conf the original config - * @return a copy with the security settings added - */ - private static Configuration addSecurityConfiguration(Configuration conf) { - // Make a copy so we don't mutate it. Also use an YarnConfiguration to - // force loading of yarn-site.xml. - conf = new YarnConfiguration(conf); - conf.set(CommonConfigurationKeys.HADOOP_SECURITY_SERVICE_USER_NAME_KEY, - conf.get(YarnConfiguration.RM_PRINCIPAL, "")); - return conf; + return opts; } - @Override public int run(String[] args) throws Exception { + + int exitCode = -1; + + AdminCommandHandler adminCmdHandler = getAdminCommandHandler(); + ClientCommandHandler clientCmdHandler = getClientCommandHandler(); + + // Build options + Options opts = buildOptions(adminCmdHandler, clientCmdHandler); + if (args.length < 1) { - printUsage(); + printUsage(false, adminCmdHandler, clientCmdHandler); return -1; } - int exitCode = -1; - int i = 0; - String cmd = args[i++]; + // Handle command separate + if (handleHelpCommand(args, adminCmdHandler, clientCmdHandler)) { + return 0; + } - if ("-help".equals(cmd)) { - exitCode = 0; - if (args.length >= 2) { - printHelpMsg(args[i]); + CommandLine cliParser; + CommandHandler handler = null; + try { + cliParser = new GnuParser().parse(opts, args); + handler = adminCmdHandler.canHandleCommand(cliParser) ? + adminCmdHandler : + clientCmdHandler.canHandleCommand(cliParser) ? + clientCmdHandler : + null; + if (handler == null) { + errOut.println(INVALID_COMMAND_USAGE); + printUsage(false, adminCmdHandler, clientCmdHandler); + return exitCode; } else { - printHelpMsg(""); + return handler.handleCommand(cliParser); } + } catch (UnrecognizedOptionException e) { + errOut.println(INVALID_COMMAND_USAGE); + printUsage(false, adminCmdHandler, clientCmdHandler); + return exitCode; + } catch (MissingArgumentException ex) { + errOut.println(MISSING_ARGUMENT); + printUsage(true, adminCmdHandler, clientCmdHandler); + return exitCode; + } catch (IllegalArgumentException arge) { + errOut.println(arge.getLocalizedMessage()); + // print admin command detail + printUsage(true, handler); + return exitCode; + } catch (Exception e) { + errOut.println(e.toString()); + printUsage(true, handler); return exitCode; } + } - try { - if ("-replace".equals(cmd)) { - exitCode = handleNodeAttributeMapping(args, - AttributeMappingOperationType.REPLACE); - } else if ("-add".equals(cmd)) { - exitCode = - handleNodeAttributeMapping(args, AttributeMappingOperationType.ADD); - } else if ("-remove".equals(cmd)) { - exitCode = handleNodeAttributeMapping(args, - AttributeMappingOperationType.REMOVE); + private boolean handleHelpCommand(String[] args, CommandHandler... handlers) + throws UnsupportedEncodingException { + if (args[0].equals("-help")) { + if (args.length == 2) { + printUsage(args[1], true, handlers); } else { - exitCode = -1; - errOut.println(cmd.substring(1) + ": Unknown command"); - printUsage(); + printUsage(true, handlers); } - } catch (IllegalArgumentException arge) { - exitCode = -1; - errOut.println(cmd.substring(1) + ": " + arge.getLocalizedMessage()); - printUsage(cmd); - } catch (RemoteException e) { - // - // This is a error returned by hadoop server. Print - // out the first line of the error message, ignore the stack trace. - exitCode = -1; - try { - String[] content; - content = e.getLocalizedMessage().split("\n"); - errOut.println(cmd.substring(1) + ": " + content[0]); - } catch (Exception ex) { - errOut.println(cmd.substring(1) + ": " + ex.getLocalizedMessage()); + return true; + } + return false; + } + + public static void main(String[] args) throws Exception { + int result = ToolRunner.run(new NodeAttributesCLI(), args); + System.exit(result); + } + + /** + * Abstract class for command handler. + */ + public static abstract class CommandHandler extends Configured { + + private Options options; + + private LinkedList<String> order = new LinkedList<>(); + private String header; + + protected CommandHandler(String header) { + this(new YarnConfiguration()); + this.header = header; + } + + protected CommandHandler(Configuration conf) { + super(conf); + options = buildOptions(); + } + + public boolean canHandleCommand(CommandLine parse) { + ArrayList<Option> arrayList = new ArrayList<Option>(options.getOptions()); + return arrayList.stream().anyMatch(opt -> parse.hasOption(opt.getOpt())); + } + + public abstract int handleCommand(CommandLine parse) + throws IOException, YarnException; + + public abstract Options buildOptions(); + + public Options getOptions() { + return options; + } + + public boolean getHelp(String cmd, StringBuilder strcnd, boolean addDesc) { + Option opt = options.getOption(cmd); + if (opt != null) { + strcnd.append(DEFAULT_SEPARATOR).append(" -").append(opt.getOpt()); + if (opt.hasArg()) { + strcnd.append(" <").append(opt.getArgName()).append(">"); + } + if (addDesc) { + strcnd.append(DEFAULT_SEPARATOR).append("\t") + .append(opt.getDescription()); + } } - } catch (Exception e) { - exitCode = -1; - errOut.println(cmd.substring(1) + ": " + e.getLocalizedMessage()); + return opt == null; + } + + public void getHelp(StringBuilder builder, boolean description) { + builder.append(DEFAULT_SEPARATOR).append(DEFAULT_SEPARATOR) + .append(header); + for (String option : order) { + getHelp(option, builder, description); + } + } + + protected void addOrder(String key){ + order.add(key); } - return exitCode; } - private int handleNodeAttributeMapping(String args[], - AttributeMappingOperationType operation) - throws IOException, YarnException, ParseException { - Options opts = new Options(); - opts.addOption(operation.name().toLowerCase(), true, - operation.name().toLowerCase()); - opts.addOption("failOnUnknownNodes", false, "Fail on unknown nodes."); - int exitCode = -1; - CommandLine cliParser = null; - try { - cliParser = new GnuParser().parse(opts, args); - } catch (MissingArgumentException ex) { - errOut.println(NO_MAPPING_ERR_MSG); - printUsage(args[0]); - return exitCode; + /** + * Client commands handler. + */ + public static class ClientCommandHandler extends CommandHandler { + + private static final String LIST_ALL_ATTRS = "list"; + + private static final String NODESTOATTR = "nodestoattributes"; + private static final String NODES = "nodes"; + + private static final String ATTRTONODES = "attributestonodes"; + private static final String ATTRIBUTES = "attributes"; + + public static final String SPLITPATTERN = "/"; + + private static final String NODEATTRIBUTE = + "%40s\t%10s\t%20s" + DEFAULT_SEPARATOR; + private static final String NODEATTRIBUTEINFO = + "%40s\t%15s" + DEFAULT_SEPARATOR; + private static final String HOSTNAMEVAL = "%40s\t%15s" + DEFAULT_SEPARATOR; + + private PrintStream sysOut = System.out; + + public ClientCommandHandler() { + super("Client Commands:"); + + } + + public void setSysOut(PrintStream out) { + this.sysOut = out; + } + + @Override + public int handleCommand(CommandLine parse) + throws IOException, YarnException { + if (parse.hasOption(LIST_ALL_ATTRS)) { + return printClusterAttributes(); + } else if (parse.hasOption(NODESTOATTR)) { + String[] nodes = new String[0]; + if (parse.hasOption(NODES)) { + nodes = parse.getOptionValues(NODES); + } + return printAttributesByNode(nodes); + } else if (parse.hasOption(ATTRTONODES)) { + String[] attrKeys = {}; + if (parse.hasOption(ATTRIBUTES)) { + attrKeys = parse.getOptionValues(ATTRIBUTES); + } + return printNodesByAttributes(attrKeys); + } + return 0; + } + + protected ApplicationClientProtocol createApplicationProtocol() + throws IOException { + // Get the current configuration + final YarnConfiguration conf = new YarnConfiguration(getConf()); + return ClientRMProxy.createRMProxy(conf, ApplicationClientProtocol.class); + } + + public int printNodesByAttributes(String[] attrs) + throws YarnException, IOException { + ApplicationClientProtocol protocol = createApplicationProtocol(); + HashSet<NodeAttributeKey> set = new HashSet<>(); + + for (String attr : attrs) { + String[] attrFields = attr.split(SPLITPATTERN); + if (attrFields.length == 1) { + set.add(NodeAttributeKey.newInstance(attrFields[0])); + } else if (attrFields.length == 2) { + set.add(NodeAttributeKey.newInstance(attrFields[0], attrFields[1])); + } else { + throw new IllegalArgumentException( + " Attribute format not correct. Should be <[prefix]/[name]> :" + + attr); + } + } + + GetAttributesToNodesRequest request = + GetAttributesToNodesRequest.newInstance(set); + GetAttributesToNodesResponse response = + protocol.getAttributesToNodes(request); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter writer = new PrintWriter( + new OutputStreamWriter(baos, Charset.forName("UTF-8"))); + writer.format(HOSTNAMEVAL, "Hostname", "Attribute-value"); + response.getAttributesToNodes().forEach((attributeKey, v) -> { + writer.println(getKeyString(attributeKey) + " :"); + v.iterator().forEachRemaining(attrVal -> writer + .format(HOSTNAMEVAL, attrVal.getHostname(), + attrVal.getAttributeValue())); + }); + writer.close(); + sysOut.println(baos.toString("UTF-8")); + return 0; + } + + private int printAttributesByNode(String[] nodeArray) + throws YarnException, IOException { + ApplicationClientProtocol protocol = createApplicationProtocol(); + HashSet<String> nodes = new HashSet<>(Arrays.asList(nodeArray)); + GetNodesToAttributesRequest request = + GetNodesToAttributesRequest.newInstance(nodes); + GetNodesToAttributesResponse response = + protocol.getNodesToAttributes(request); + Map<String, Set<NodeAttribute>> nodeToAttrs = + response.getNodeToAttributes(); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter writer = new PrintWriter( + new OutputStreamWriter(baos, Charset.forName("UTF-8"))); + writer.printf(NODEATTRIBUTE, "Attribute", "Type", "Value"); + nodeToAttrs.forEach((node, v) -> { + // print node header + writer.println(node + ":"); + v.iterator().forEachRemaining(attr -> writer + .format(NODEATTRIBUTE, getKeyString(attr.getAttributeKey()), + attr.getAttributeType().name(), attr.getAttributeValue())); + }); + writer.close(); + sysOut.println(baos.toString("UTF-8")); + return 0; + } + + private int printClusterAttributes() throws IOException, YarnException { + ApplicationClientProtocol protocol = createApplicationProtocol(); + GetClusterNodeAttributesRequest request = + GetClusterNodeAttributesRequest.newInstance(); + GetClusterNodeAttributesResponse response = + protocol.getClusterNodeAttributes(request); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter writer = new PrintWriter( + new OutputStreamWriter(baos, Charset.forName("UTF-8"))); + writer.format(NODEATTRIBUTEINFO, "Attribute", "Type"); + for (NodeAttributeInfo attr : response.getNodeAttributes()) { + writer.format(NODEATTRIBUTEINFO, getKeyString(attr.getAttributeKey()), + attr.getAttributeType().name()); + } + writer.close(); + sysOut.println(baos.toString("UTF-8")); + return 0; + } + + private String getKeyString(NodeAttributeKey key) { + StringBuilder sb = new StringBuilder(); + sb.append(key.getAttributePrefix()).append("/") + .append(key.getAttributeName()); + return sb.toString(); + } + + @Override + public Options buildOptions() { + Options clientOptions = new Options(); + clientOptions.addOption( + new Option(LIST_ALL_ATTRS, false, "List all attributes in cluster")); + + // group by command + OptionGroup nodeToAttr = new OptionGroup(); + Option attrtonodes = new Option(NODESTOATTR, false, + "Lists all mapping to nodes to attributes"); + Option nodes = new Option(NODES, + "Works with [" + LIST_ALL_ATTRS + "] to specify node hostnames " + + "whose mappings are required to be displayed."); + nodes.setValueSeparator(','); + nodes.setArgName("Host Names"); + nodes.setArgs(Option.UNLIMITED_VALUES); + nodeToAttr.addOption(attrtonodes); + nodeToAttr.addOption(nodes); + clientOptions.addOptionGroup(nodeToAttr); + + // Defines as groups to add extendability for later + OptionGroup attrToNodes = new OptionGroup(); + attrToNodes.addOption(new Option(ATTRTONODES, false, + "Displays mapping of " + + "attributes to nodes and attribute values grouped by " + + "attributes")); + Option attrs = new Option(ATTRIBUTES, "Works with [" + ATTRTONODES + + "] to specify attributes whose mapping " + + "are required to be displayed."); + attrs.setValueSeparator(','); + attrs.setArgName("Attributes"); + attrs.setArgs(Option.UNLIMITED_VALUES); + attrToNodes.addOption(attrs); + clientOptions.addOptionGroup(attrToNodes); + + // DEFINE ORDER + addOrder(LIST_ALL_ATTRS); + addOrder(NODESTOATTR); + addOrder(NODES); + addOrder(ATTRTONODES); + addOrder(ATTRIBUTES); + return clientOptions; } - List<NodeToAttributes> buildNodeLabelsMapFromStr = - buildNodeLabelsMapFromStr( - cliParser.getOptionValue(operation.name().toLowerCase()), - operation != AttributeMappingOperationType.REPLACE, operation); - NodesToAttributesMappingRequest request = NodesToAttributesMappingRequest - .newInstance(operation, buildNodeLabelsMapFromStr, - cliParser.hasOption("failOnUnknownNodes")); - ResourceManagerAdministrationProtocol adminProtocol = createAdminProtocol(); - adminProtocol.mapAttributesToNodes(request); - return 0; } /** - * args are expected to be of the format - * node1:java(string)=8,ssd(boolean)=false node2:ssd(boolean)=true + * Admin commands handler. */ - private List<NodeToAttributes> buildNodeLabelsMapFromStr(String args, - boolean validateForAttributes, AttributeMappingOperationType operation) { - Map<String,NodeToAttributes> nodeToAttributesMap = new HashMap<>(); - for (String nodeToAttributesStr : args.split("[ \n]")) { - // for each node to attribute mapping - nodeToAttributesStr = nodeToAttributesStr.trim(); - if (nodeToAttributesStr.isEmpty() - || nodeToAttributesStr.startsWith("#")) { - continue; + public static class AdminCommandHandler extends CommandHandler { + + private static final String ADD = "add"; + private static final String REMOVE = "remove"; + private static final String REPLACE = "replace"; + private static final String FAILUNKNOWNNODES = "failOnUnknownNodes"; + + AdminCommandHandler() { + super("Admin Commands:"); + } + + @Override + public Options buildOptions() { + Options adminOptions = new Options(); + Option replace = new Option(REPLACE, true, + "Replace the node to attributes mapping information at the" + + " ResourceManager with the new mapping. Currently" + + " supported attribute type. And string is the default" + + " type too. Attribute value if not specified for string" + + " type value will be considered as empty string." + + " Replaced node-attributes should not violate the" + + " existing attribute to attribute type mapping."); + replace.setArgName("\"node1:attribute[(type)][=value],attribute1[=value]," + + "attribute2 node2:attribute2[=value],attribute3\""); + replace.setArgs(1); + adminOptions.addOption(replace); + + Option add = new Option(ADD, true, + "Adds or updates the node to attributes mapping information" + + " at the ResourceManager. Currently supported attribute" + + " type is string. And string is the default type too." + + " Attribute value if not specified for string type" + + " value will be considered as empty string. Added or" + + " updated node-attributes should not violate the" + + " existing attribute to attribute type mapping."); + add.setArgName("\"node1:attribute[(type)][=value],attribute1[=value]," + + "attribute2 node2:attribute2[=value],attribute3\""); + add.setArgs(1); + adminOptions.addOption(add); + + Option remove = new Option(REMOVE, true, + "Removes the specified node to attributes mapping" + + " information at the ResourceManager"); + remove.setArgName("\"node1:attribute,attribute1 node2:attribute2\""); + remove.setArgs(1); + adminOptions.addOption(remove); + + adminOptions.addOption(new Option(FAILUNKNOWNNODES, false, + "Can be used optionally along with [add,remove,replace] options. " + + "When set, command will fail if specified nodes are unknown.")); + + // DEFINE ORDER + addOrder(REPLACE); + addOrder(ADD); + addOrder(REMOVE); + addOrder(FAILUNKNOWNNODES); + + return adminOptions; + } + + protected ResourceManagerAdministrationProtocol createAdminProtocol() + throws IOException { + // Get the current configuration + final YarnConfiguration conf = new YarnConfiguration(getConf()); + return ClientRMProxy + .createRMProxy(conf, ResourceManagerAdministrationProtocol.class); + } + + public int handleCommand(CommandLine cliParser) + throws IOException, YarnException { + String operation = null; + if (cliParser.hasOption(ADD)) { + operation = ADD; + } else if (cliParser.hasOption(REMOVE)) { + operation = REMOVE; + } else if (cliParser.hasOption(REPLACE)) { + operation = REPLACE; } - if (nodeToAttributesStr.indexOf(":") == -1) { + if (operation != null) { + List<NodeToAttributes> buildNodeLabelsListFromStr = + buildNodeLabelsListFromStr(cliParser.getOptionValue(operation), + !operation.equals(REPLACE), operation); + NodesToAttributesMappingRequest request = + NodesToAttributesMappingRequest.newInstance( + AttributeMappingOperationType.valueOf(operation.toUpperCase()), + buildNodeLabelsListFromStr, + cliParser.hasOption(FAILUNKNOWNNODES)); + ResourceManagerAdministrationProtocol adminProtocol = + createAdminProtocol(); + adminProtocol.mapAttributesToNodes(request); + } else { + // Handle case for only failOnUnknownNodes passed throw new IllegalArgumentException( - INVALID_MAPPING_ERR_MSG + nodeToAttributesStr); + getOptions().getOption(FAILUNKNOWNNODES).getDescription()); } - String[] nodeToAttributes = nodeToAttributesStr.split(":"); - Preconditions.checkArgument(!nodeToAttributes[0].trim().isEmpty(), - "Node name cannot be empty"); - String node = nodeToAttributes[0]; - String[] attributeNameValueType = null; - List<NodeAttribute> attributesList = new ArrayList<>(); - NodeAttributeType attributeType = NodeAttributeType.STRING; - String attributeValue; - String attributeName; - Set<String> attributeNamesMapped = new HashSet<>(); - - String attributesStr[]; - if (nodeToAttributes.length == 2) { - // fetching multiple attributes for a node - attributesStr = nodeToAttributes[1].split(","); - for (String attributeStr : attributesStr) { - // get information about each attribute. - attributeNameValueType = attributeStr.split("="); // to find name - // value - Preconditions.checkArgument( - !(attributeNameValueType[0] == null - || attributeNameValueType[0].isEmpty()), - "Attribute name cannot be null or empty"); - attributeValue = attributeNameValueType.length > 1 - ? attributeNameValueType[1] : ""; - int indexOfOpenBracket = attributeNameValueType[0].indexOf("("); - if (indexOfOpenBracket == -1) { - attributeName = attributeNameValueType[0]; - } else if (indexOfOpenBracket == 0) { - throw new IllegalArgumentException("Attribute for node " + node - + " is not properly configured : " + attributeStr); - } else { - // attribute type has been explicitly configured - int indexOfCloseBracket = attributeNameValueType[0].indexOf(")"); - if (indexOfCloseBracket == -1 - || indexOfCloseBracket < indexOfOpenBracket) { + return 0; + } + + /** + * args are expected to be of the format + * node1:java(string)=8,ssd(boolean)=false node2:ssd(boolean)=true. + */ + private List<NodeToAttributes> buildNodeLabelsListFromStr(String args, + boolean validateForAttributes, String operation) { + Map<String, NodeToAttributes> nodeToAttributesMap = new HashMap<>(); + for (String nodeToAttributesStr : args.split("[ \n]")) { + // for each node to attribute mapping + nodeToAttributesStr = nodeToAttributesStr.trim(); + if (nodeToAttributesStr.isEmpty() || nodeToAttributesStr + .startsWith("#")) { + continue; + } + if (nodeToAttributesStr.indexOf(":") == -1) { + throw new IllegalArgumentException( + INVALID_MAPPING_ERR_MSG + nodeToAttributesStr); + } + String[] nodeToAttributes = nodeToAttributesStr.split(":"); + Preconditions.checkArgument(!nodeToAttributes[0].trim().isEmpty(), + "Node name cannot be empty"); + String node = nodeToAttributes[0]; + String[] attributeNameValueType = null; + List<NodeAttribute> attributesList = new ArrayList<>(); + NodeAttributeType attributeType = NodeAttributeType.STRING; + String attributeValue; + String attributeName; + Set<String> attributeNamesMapped = new HashSet<>(); + + String[] attributesStr; + if (nodeToAttributes.length == 2) { + // fetching multiple attributes for a node + attributesStr = nodeToAttributes[1].split(","); + for (String attributeStr : attributesStr) { + // get information about each attribute. + attributeNameValueType = attributeStr.split("="); // to find name + // value + Preconditions.checkArgument( + !(attributeNameValueType[0] == null || attributeNameValueType[0] + .isEmpty()), "Attribute name cannot be null or empty"); + attributeValue = attributeNameValueType.length > 1 ? + attributeNameValueType[1] : + ""; + int indexOfOpenBracket = attributeNameValueType[0].indexOf("("); + if (indexOfOpenBracket == -1) { + attributeName = attributeNameValueType[0]; + } else if (indexOfOpenBracket == 0) { throw new IllegalArgumentException("Attribute for node " + node - + " is not properly Configured : " + attributeStr); + + " is not properly configured : " + attributeStr); + } else { + // attribute type has been explicitly configured + int indexOfCloseBracket = attributeNameValueType[0].indexOf(")"); + if (indexOfCloseBracket == -1 + || indexOfCloseBracket < indexOfOpenBracket) { + throw new IllegalArgumentException("Attribute for node " + node + + " is not properly Configured : " + attributeStr); + } + String attributeTypeStr; + attributeName = + attributeNameValueType[0].substring(0, indexOfOpenBracket); + attributeTypeStr = attributeNameValueType[0] + .substring(indexOfOpenBracket + 1, indexOfCloseBracket); + try { + attributeType = NodeAttributeType + .valueOf(attributeTypeStr.trim().toUpperCase()); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + "Invalid Attribute type configuration : " + attributeTypeStr + + " in " + attributeStr); + } } - String attributeTypeStr; - attributeName = - attributeNameValueType[0].substring(0, indexOfOpenBracket); - attributeTypeStr = attributeNameValueType[0] - .substring(indexOfOpenBracket + 1, indexOfCloseBracket); - try { - attributeType = NodeAttributeType - .valueOf(attributeTypeStr.trim().toUpperCase()); - } catch (IllegalArgumentException e) { - throw new IllegalArgumentException( - "Invalid Attribute type configuration : " + attributeTypeStr - + " in " + attributeStr); + if (attributeNamesMapped.contains(attributeName)) { + throw new IllegalArgumentException("Attribute " + attributeName + + " has been mapped more than once in : " + + nodeToAttributesStr); } + // TODO when we support different type of attribute type we need to + // cross verify whether input attributes itself is not violating + // attribute Name to Type mapping. + attributesList.add(NodeAttribute + .newInstance(NodeAttribute.PREFIX_CENTRALIZED, + attributeName.trim(), attributeType, + attributeValue.trim())); } - if (attributeNamesMapped.contains(attributeName)) { - throw new IllegalArgumentException("Attribute " + attributeName - + " has been mapped more than once in : " - + nodeToAttributesStr); - } - // TODO when we support different type of attribute type we need to - // cross verify whether input attributes itself is not violating - // attribute Name to Type mapping. - attributesList - .add(NodeAttribute.newInstance(NodeAttribute.PREFIX_CENTRALIZED, - attributeName.trim(), attributeType, attributeValue.trim())); } + if (validateForAttributes) { + Preconditions.checkArgument((attributesList.size() > 0), + "Attributes cannot be null or empty for Operation [" + operation + + "] on the node " + node); + } + nodeToAttributesMap + .put(node, NodeToAttributes.newInstance(node, attributesList)); } - if (validateForAttributes) { - Preconditions.checkArgument((attributesList.size() > 0), - "Attributes cannot be null or empty for Operation " - + operation.name() + " on the node " + node); + + if (nodeToAttributesMap.isEmpty()) { + throw new IllegalArgumentException(NO_MAPPING_ERR_MSG); } - nodeToAttributesMap - .put(node,NodeToAttributes.newInstance(node, attributesList)); + return Lists.newArrayList(nodeToAttributesMap.values()); } - if (nodeToAttributesMap.isEmpty()) { - throw new IllegalArgumentException(NO_MAPPING_ERR_MSG); + @Override + public void setConf(Configuration conf) { + if (conf != null) { + conf = addSecurityConfiguration(conf); + } + super.setConf(conf); + } + + /** + * Add the requisite security principal settings to the given Configuration, + * returning a copy. + * + * @param conf the original config + * @return a copy with the security settings added + */ + private Configuration addSecurityConfiguration(Configuration conf) { + // Make a copy so we don't mutate it. Also use an YarnConfiguration to + // force loading of yarn-site.xml. + conf = new YarnConfiguration(conf); + conf.set(CommonConfigurationKeys.HADOOP_SECURITY_SERVICE_USER_NAME_KEY, + conf.get(YarnConfiguration.RM_PRINCIPAL, "")); + return conf; } - return Lists.newArrayList(nodeToAttributesMap.values()); - } - public static void main(String[] args) throws Exception { - int result = ToolRunner.run(new NodeAttributesCLI(), args); - System.exit(result); } } http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeCLI.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeCLI.java index e9253eb..44e9870 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeCLI.java @@ -44,7 +44,6 @@ import org.apache.hadoop.yarn.api.records.NodeId; import org.apache.hadoop.yarn.api.records.NodeReport; import org.apache.hadoop.yarn.api.records.NodeState; import org.apache.hadoop.yarn.exceptions.YarnException; -import org.apache.hadoop.yarn.util.ConverterUtils; @Private @Unstable @@ -307,6 +306,18 @@ public class NodeCLI extends YarnCLI { Collections.sort(nodeLabelsList); nodeReportStr.println(StringUtils.join(nodeLabelsList.iterator(), ',')); + if (nodeReport.getNodeAttributes().size() > 0) { + ArrayList nodeAtrs = new ArrayList<>(nodeReport.getNodeAttributes()); + nodeReportStr.print("\tNode Attributes : "); + nodeReportStr.println(nodeAtrs.get(0).toString()); + for (int index = 1; index < nodeAtrs.size(); index++) { + nodeReportStr.println( + String.format("\t%18s%s", "", nodeAtrs.get(index).toString())); + } + } else { + nodeReportStr.println("\tNode Attributes : "); + } + nodeReportStr.print("\tResource Utilization by Node : "); if (nodeReport.getNodeUtilization() != null) { nodeReportStr.print("PMem:" http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java index 5a0f049..26afe6f 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java @@ -18,6 +18,9 @@ package org.apache.hadoop.yarn.client.cli; +import org.apache.hadoop.yarn.api.records.NodeAttributeInfo; +import org.apache.hadoop.yarn.api.records.NodeAttributeKey; +import org.apache.hadoop.yarn.api.records.NodeAttributeType; import static org.junit.Assert.assertEquals; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; @@ -74,7 +77,32 @@ public class TestClusterCLI { pw.close(); verify(sysOut).println(baos.toString("UTF-8")); } - + + @Test + public void testGetClusterNodeAttributes() throws Exception { + YarnClient client = mock(YarnClient.class); + when(client.getClusterAttributes()).thenReturn(ImmutableSet + .of(NodeAttributeInfo.newInstance(NodeAttributeKey.newInstance("GPU"), + NodeAttributeType.STRING), NodeAttributeInfo + .newInstance(NodeAttributeKey.newInstance("CPU"), + NodeAttributeType.STRING))); + ClusterCLI cli = new ClusterCLI(); + cli.setClient(client); + cli.setSysOutPrintStream(sysOut); + cli.setSysErrPrintStream(sysErr); + + int rc = cli.run(new String[] {ClusterCLI.CMD, + "-" + ClusterCLI.LIST_CLUSTER_ATTRIBUTES}); + assertEquals(0, rc); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter pw = new PrintWriter(baos); + pw.println("rm.yarn.io/GPU(STRING)"); + pw.println("rm.yarn.io/CPU(STRING)"); + pw.close(); + verify(sysOut).println(baos.toString("UTF-8")); + } + @Test public void testGetClusterNodeLabelsWithLocalAccess() throws Exception { YarnClient client = mock(YarnClient.class); @@ -157,6 +185,8 @@ public class TestClusterCLI { pw.println(" option is UNSTABLE, could be"); pw.println(" removed in future releases."); pw.println(" -h,--help Displays help for all commands."); + pw.println(" -lna,--list-node-attributes List cluster node-attribute"); + pw.println(" collection"); pw.println(" -lnl,--list-node-labels List cluster node-label"); pw.println(" collection"); pw.close(); http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java index bbd5ca3..7f48493 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java @@ -18,6 +18,20 @@ package org.apache.hadoop.yarn.client.cli; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.hadoop.yarn.api.ApplicationClientProtocol; +import org.apache.hadoop.yarn.api.protocolrecords.GetAttributesToNodesRequest; +import org.apache.hadoop.yarn.api.protocolrecords.GetAttributesToNodesResponse; +import org.apache.hadoop.yarn.api.protocolrecords.GetClusterNodeAttributesRequest; +import org.apache.hadoop.yarn.api.protocolrecords.GetClusterNodeAttributesResponse; + +import org.apache.hadoop.yarn.api.protocolrecords.GetNodesToAttributesRequest; +import org.apache.hadoop.yarn.api.protocolrecords.GetNodesToAttributesResponse; +import org.apache.hadoop.yarn.api.records.NodeAttributeInfo; +import org.apache.hadoop.yarn.api.records.NodeAttributeKey; +import org.apache.hadoop.yarn.api.records.NodeToAttributeValue; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.Matchers.any; @@ -29,8 +43,8 @@ import java.io.IOException; import java.io.PrintStream; import java.util.ArrayList; import java.util.List; +import java.util.Set; -import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.yarn.api.records.NodeAttribute; import org.apache.hadoop.yarn.api.records.NodeAttributeType; import org.apache.hadoop.yarn.exceptions.YarnException; @@ -56,75 +70,122 @@ public class TestNodeAttributesCLI { private static final Logger LOG = LoggerFactory.getLogger(TestNodeAttributesCLI.class); private ResourceManagerAdministrationProtocol admin; - private NodesToAttributesMappingRequest request; + private ApplicationClientProtocol client; + private NodesToAttributesMappingRequest nodeToAttrRequest; private NodeAttributesCLI nodeAttributesCLI; private ByteArrayOutputStream errOutBytes = new ByteArrayOutputStream(); + private ByteArrayOutputStream sysOutBytes = new ByteArrayOutputStream(); private String errOutput; + private String sysOutput; @Before public void configure() throws IOException, YarnException { + admin = mock(ResourceManagerAdministrationProtocol.class); + client = mock(ApplicationClientProtocol.class); when(admin.mapAttributesToNodes(any(NodesToAttributesMappingRequest.class))) .thenAnswer(new Answer<NodesToAttributesMappingResponse>() { @Override public NodesToAttributesMappingResponse answer( InvocationOnMock invocation) throws Throwable { - request = + nodeToAttrRequest = (NodesToAttributesMappingRequest) invocation.getArguments()[0]; return NodesToAttributesMappingResponse.newInstance(); } }); - nodeAttributesCLI = new NodeAttributesCLI(new Configuration()) { + nodeAttributesCLI = new NodeAttributesCLI() { @Override - protected ResourceManagerAdministrationProtocol createAdminProtocol() - throws IOException { - return admin; + protected AdminCommandHandler getAdminCommandHandler() { + return new AdminCommandHandler() { + @Override + protected ResourceManagerAdministrationProtocol createAdminProtocol() + throws IOException { + return admin; + } + }; } - }; + @Override + protected ClientCommandHandler getClientCommandHandler() { + ClientCommandHandler handler = new ClientCommandHandler() { + @Override + protected ApplicationClientProtocol createApplicationProtocol() + throws IOException { + return client; + } + }; + handler.setSysOut(new PrintStream(sysOutBytes)); + return handler; + } + }; nodeAttributesCLI.setErrOut(new PrintStream(errOutBytes)); } @Test public void testHelp() throws Exception { - String[] args = new String[] { "-help", "-replace" }; + String[] args = new String[] {"-help", "-replace"}; assertTrue("It should have succeeded help for replace", 0 == runTool(args)); - assertOutputContains( - "-replace <\"node1:attribute[(type)][=value],attribute1" - + "[=value],attribute2 node2:attribute2[=value],attribute3\"> :"); - assertOutputContains("Replace the node to attributes mapping information at" + assertErrorContains("-replace <\"node1:attribute[(type)][=value],attribute1" + + "[=value],attribute2 node2:attribute2[=value],attribute3\">"); + assertErrorContains("Replace the node to attributes mapping information at" + " the ResourceManager with the new mapping. Currently supported" + " attribute type. And string is the default type too. Attribute value" + " if not specified for string type value will be considered as empty" + " string. Replaced node-attributes should not violate the existing" + " attribute to attribute type mapping."); - args = new String[] { "-help", "-remove" }; + args = new String[] {"-help", "-remove"}; assertTrue("It should have succeeded help for replace", 0 == runTool(args)); - assertOutputContains( - "-remove <\"node1:attribute,attribute1" + " node2:attribute2\"> :"); - assertOutputContains("Removes the specified node to attributes mapping" + assertErrorContains( + "-remove <\"node1:attribute,attribute1" + " node2:attribute2\">"); + assertErrorContains("Removes the specified node to attributes mapping" + " information at the ResourceManager"); - args = new String[] { "-help", "-add" }; + args = new String[] {"-help", "-add"}; assertTrue("It should have succeeded help for replace", 0 == runTool(args)); - assertOutputContains("-add <\"node1:attribute[(type)][=value]," - + "attribute1[=value],attribute2 node2:attribute2[=value],attribute3\">" - + " :"); - assertOutputContains("Adds or updates the node to attributes mapping" + assertErrorContains("-add <\"node1:attribute[(type)][=value]," + + "attribute1[=value],attribute2 node2:attribute2[=value]," + + "attribute3\">"); + assertErrorContains("Adds or updates the node to attributes mapping" + " information at the ResourceManager. Currently supported attribute" + " type is string. And string is the default type too. Attribute value" + " if not specified for string type value will be considered as empty" + " string. Added or updated node-attributes should not violate the" + " existing attribute to attribute type mapping."); - args = new String[] { "-help", "-failOnUnknownNodes" }; + args = new String[] {"-help", "-failOnUnknownNodes"}; assertTrue("It should have succeeded help for replace", 0 == runTool(args)); - assertOutputContains("-failOnUnknownNodes :"); - assertOutputContains("Can be used optionally along with other options. When" - + " its set, it will fail if specified nodes are unknown."); + assertErrorContains("-failOnUnknownNodes"); + assertErrorContains("Can be used optionally along with [add,remove," + + "replace] options. When set, command will fail if specified nodes " + + "are unknown."); + + args = new String[] {"-help", "-list"}; + assertTrue("It should have succeeded help for replace", 0 == runTool(args)); + assertErrorContains("-list"); + assertErrorContains("List all attributes in cluster"); + + args = new String[] {"-help", "-nodes"}; + assertTrue("It should have succeeded help for replace", 0 == runTool(args)); + assertErrorContains("-nodes"); + assertErrorContains( + "Works with [list] to specify node hostnames whose mappings " + + "are required to be displayed."); + + args = new String[] {"-help", "-attributes"}; + assertTrue("It should have succeeded help for replace", 0 == runTool(args)); + assertErrorContains("-attributes"); + assertErrorContains( + "Works with [attributestonodes] to specify attributes whose mapping " + + "are required to be displayed."); + + args = new String[] {"-help", "-attributestonodes"}; + assertTrue("It should have succeeded help for replace", 0 == runTool(args)); + assertErrorContains("-attributestonodes"); + assertErrorContains("Displays mapping of attributes to nodes and attribute " + + "values grouped by attributes"); } @Test @@ -133,62 +194,62 @@ public class TestNodeAttributesCLI { // failure scenarios // -------------------------------- // parenthesis not match - String[] args = new String[] { "-replace", "x(" }; + String[] args = new String[] {"-replace", "x("}; assertTrue("It should have failed as no node is specified", 0 != runTool(args)); assertFailureMessageContains(NodeAttributesCLI.INVALID_MAPPING_ERR_MSG); // parenthesis not match - args = new String[] { "-replace", "x:(=abc" }; + args = new String[] {"-replace", "x:(=abc"}; assertTrue( "It should have failed as no closing parenthesis is not specified", 0 != runTool(args)); assertFailureMessageContains( "Attribute for node x is not properly configured : (=abc"); - args = new String[] { "-replace", "x:()=abc" }; + args = new String[] {"-replace", "x:()=abc"}; assertTrue("It should have failed as no type specified inside parenthesis", 0 != runTool(args)); assertFailureMessageContains( "Attribute for node x is not properly configured : ()=abc"); - args = new String[] { "-replace", ":x(string)" }; + args = new String[] {"-replace", ":x(string)"}; assertTrue("It should have failed as no node is specified", 0 != runTool(args)); assertFailureMessageContains("Node name cannot be empty"); // Not expected key=value specifying inner parenthesis - args = new String[] { "-replace", "x:(key=value)" }; + args = new String[] {"-replace", "x:(key=value)"}; assertTrue(0 != runTool(args)); assertFailureMessageContains( "Attribute for node x is not properly configured : (key=value)"); // Should fail as no attributes specified - args = new String[] { "-replace" }; + args = new String[] {"-replace"}; assertTrue("Should fail as no attribute mappings specified", 0 != runTool(args)); - assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG); + assertFailureMessageContains(NodeAttributesCLI.MISSING_ARGUMENT); // no labels, should fail - args = new String[] { "-replace", "-failOnUnknownNodes", - "x:key(string)=value,key2=val2" }; + args = new String[] {"-replace", "-failOnUnknownNodes", + "x:key(string)=value,key2=val2"}; assertTrue("Should fail as no attribute mappings specified for replace", 0 != runTool(args)); - assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG); + assertFailureMessageContains(NodeAttributesCLI.MISSING_ARGUMENT); // no labels, should fail - args = new String[] { "-replace", " " }; + args = new String[] {"-replace", " "}; assertTrue(0 != runTool(args)); assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG); - args = new String[] { "-replace", ", " }; + args = new String[] {"-replace", ", "}; assertTrue(0 != runTool(args)); assertFailureMessageContains(NodeAttributesCLI.INVALID_MAPPING_ERR_MSG); // -------------------------------- // success scenarios // -------------------------------- - args = new String[] { "-replace", - "x:key(string)=value,key2=val2 y:key2=val23,key3 z:key4" }; + args = new String[] {"-replace", + "x:key(string)=value,key2=val2 y:key2=val23,key3 z:key4"}; assertTrue("Should not fail as attribute has been properly mapped", 0 == runTool(args)); List<NodeToAttributes> nodeAttributesList = new ArrayList<>(); @@ -221,10 +282,10 @@ public class TestNodeAttributesCLI { .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, "")); nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes)); - NodesToAttributesMappingRequest expected = - NodesToAttributesMappingRequest.newInstance( - AttributeMappingOperationType.REPLACE, nodeAttributesList, false); - assertTrue(request.equals(expected)); + NodesToAttributesMappingRequest expected = NodesToAttributesMappingRequest + .newInstance(AttributeMappingOperationType.REPLACE, nodeAttributesList, + false); + assertTrue(nodeToAttrRequest.equals(expected)); } @Test @@ -233,16 +294,17 @@ public class TestNodeAttributesCLI { // failure scenarios // -------------------------------- // parenthesis not match - String[] args = new String[] { "-remove", "x:" }; + String[] args = new String[] {"-remove", "x:"}; assertTrue("It should have failed as no node is specified", 0 != runTool(args)); assertFailureMessageContains( - "Attributes cannot be null or empty for Operation REMOVE on the node x"); + "Attributes cannot be null or empty for Operation [remove] on the " + + "node x"); // -------------------------------- // success scenarios // -------------------------------- args = - new String[] { "-remove", "x:key2,key3 z:key4", "-failOnUnknownNodes" }; + new String[] {"-remove", "x:key2,key3 z:key4", "-failOnUnknownNodes"}; assertTrue("Should not fail as attribute has been properly mapped", 0 == runTool(args)); List<NodeToAttributes> nodeAttributesList = new ArrayList<>(); @@ -259,10 +321,10 @@ public class TestNodeAttributesCLI { .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, "")); nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes)); - NodesToAttributesMappingRequest expected = - NodesToAttributesMappingRequest.newInstance( - AttributeMappingOperationType.REMOVE, nodeAttributesList, true); - assertTrue(request.equals(expected)); + NodesToAttributesMappingRequest expected = NodesToAttributesMappingRequest + .newInstance(AttributeMappingOperationType.REMOVE, nodeAttributesList, + true); + assertTrue(nodeToAttrRequest.equals(expected)); } @Test @@ -271,16 +333,16 @@ public class TestNodeAttributesCLI { // failure scenarios // -------------------------------- // parenthesis not match - String[] args = new String[] { "-add", "x:" }; + String[] args = new String[] {"-add", "x:"}; assertTrue("It should have failed as no node is specified", 0 != runTool(args)); assertFailureMessageContains( - "Attributes cannot be null or empty for Operation ADD on the node x"); + "Attributes cannot be null or empty for Operation [add] on the node x"); // -------------------------------- // success scenarios // -------------------------------- - args = new String[] { "-add", "x:key2=123,key3=abc z:key4(string)", - "-failOnUnknownNodes" }; + args = new String[] {"-add", "x:key2=123,key3=abc z:key4(string)", + "-failOnUnknownNodes"}; assertTrue("Should not fail as attribute has been properly mapped", 0 == runTool(args)); List<NodeToAttributes> nodeAttributesList = new ArrayList<>(); @@ -297,16 +359,16 @@ public class TestNodeAttributesCLI { .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, "")); nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes)); - NodesToAttributesMappingRequest expected = - NodesToAttributesMappingRequest.newInstance( - AttributeMappingOperationType.ADD, nodeAttributesList, true); - assertTrue(request.equals(expected)); + NodesToAttributesMappingRequest expected = NodesToAttributesMappingRequest + .newInstance(AttributeMappingOperationType.ADD, nodeAttributesList, + true); + assertTrue(nodeToAttrRequest.equals(expected)); // -------------------------------- // with Duplicate mappings for a host // -------------------------------- - args = new String[] { "-add", "x:key2=123,key3=abc x:key4(string)", - "-failOnUnknownNodes" }; + args = new String[] {"-add", "x:key2=123,key3=abc x:key4(string)", + "-failOnUnknownNodes"}; assertTrue("Should not fail as attribute has been properly mapped", 0 == runTool(args)); nodeAttributesList = new ArrayList<>(); @@ -315,32 +377,161 @@ public class TestNodeAttributesCLI { .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, "")); nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes)); - expected = - NodesToAttributesMappingRequest.newInstance( - AttributeMappingOperationType.ADD, nodeAttributesList, true); - assertTrue(request.equals(expected)); + expected = NodesToAttributesMappingRequest + .newInstance(AttributeMappingOperationType.ADD, nodeAttributesList, + true); + assertTrue(nodeToAttrRequest.equals(expected)); + } + + @Test + public void testListAttributes() throws Exception { + + // GetClusterNodeAttributesRequest + when(client + .getClusterNodeAttributes(any(GetClusterNodeAttributesRequest.class))) + .thenAnswer(new Answer<GetClusterNodeAttributesResponse>() { + @Override + public GetClusterNodeAttributesResponse answer( + InvocationOnMock invocation) throws Throwable { + GetClusterNodeAttributesRequest nodeAttrReq = + (GetClusterNodeAttributesRequest) invocation.getArguments()[0]; + return GetClusterNodeAttributesResponse.newInstance(ImmutableSet + .of(NodeAttributeInfo + .newInstance(NodeAttributeKey.newInstance("GPU"), + NodeAttributeType.STRING))); + } + }); + + // -------------------------------- + // Success scenarios + // -------------------------------- + String[] args = new String[] {"-list"}; + assertTrue("It should be success since it list all attributes", + 0 == runTool(args)); + assertSysOutContains("Attribute\t Type", + "rm.yarn.io/GPU\t STRING"); + } + + @Test + public void testNodeToAttributes() throws Exception { + // GetNodesToAttributesRequest response + when(client.getNodesToAttributes(any(GetNodesToAttributesRequest.class))) + .thenAnswer(new Answer<GetNodesToAttributesResponse>() { + @Override + public GetNodesToAttributesResponse answer( + InvocationOnMock invocation) throws Throwable { + GetNodesToAttributesRequest nodeToAttributes = + (GetNodesToAttributesRequest) invocation.getArguments()[0]; + return GetNodesToAttributesResponse.newInstance( + ImmutableMap.<String, Set<NodeAttribute>>builder() + .put("hostname", ImmutableSet.of(NodeAttribute + .newInstance("GPU", NodeAttributeType.STRING, "ARM"))) + .build()); + } + }); + // -------------------------------- + // Failure scenarios + // -------------------------------- + String[] args = new String[] {"-nodetoattributes", "-nodes"}; + assertTrue("It should not success since nodes are not specified", + 0 != runTool(args)); + assertErrorContains(NodeAttributesCLI.INVALID_COMMAND_USAGE); + + // Missing argument for nodes + args = new String[] {"-nodestoattributes", "-nodes"}; + assertTrue("It should not success since nodes are not specified", + 0 != runTool(args)); + assertErrorContains(NodeAttributesCLI.MISSING_ARGUMENT); + + // -------------------------------- + // Success with hostname param + // -------------------------------- + args = new String[] {"-nodestoattributes", "-nodes", "hostname"}; + assertTrue("Should return hostname to attributed list", 0 == runTool(args)); + assertSysOutContains("hostname"); + } + + @Test + public void testAttributesToNodes() throws Exception { + // GetAttributesToNodesResponse response + when(client.getAttributesToNodes(any(GetAttributesToNodesRequest.class))) + .thenAnswer(new Answer<GetAttributesToNodesResponse>() { + @Override + public GetAttributesToNodesResponse answer( + InvocationOnMock invocation) throws Throwable { + GetAttributesToNodesRequest attrToNodes = + (GetAttributesToNodesRequest) invocation.getArguments()[0]; + return GetAttributesToNodesResponse.newInstance( + ImmutableMap.<NodeAttributeKey, + List<NodeToAttributeValue>>builder() + .put(NodeAttributeKey.newInstance("GPU"), ImmutableList + .of(NodeToAttributeValue.newInstance("host1", "ARM"))) + .build()); + } + }); + // -------------------------------- + // Success scenarios + // -------------------------------- + String[] args = new String[] {"-attributestonodes"}; + assertTrue("It should be success since it list all attributes", + 0 == runTool(args)); + assertSysOutContains("Hostname\tAttribute-value", "rm.yarn.io/GPU :", + "host1\t ARM"); + + // -------------------------------- + // fail scenario argument filter missing + // -------------------------------- + args = new String[] {"-attributestonodes", "-attributes"}; + assertTrue( + "It should not success since attributes for filter are not specified", + 0 != runTool(args)); + assertErrorContains(NodeAttributesCLI.MISSING_ARGUMENT); + + // -------------------------------- + // fail scenario argument filter missing + // -------------------------------- + args = new String[] {"-attributestonodes", "-attributes", "fail/da/fail"}; + assertTrue("It should not success since attributes format is not correct", + 0 != runTool(args)); + assertErrorContains( + "Attribute format not correct. Should be <[prefix]/[name]> " + + ":fail/da/fail"); } private void assertFailureMessageContains(String... messages) { - assertOutputContains(messages); - assertOutputContains(NodeAttributesCLI.USAGE_YARN_NODE_ATTRIBUTES); + assertErrorContains(messages); + assertErrorContains(NodeAttributesCLI.USAGE_YARN_NODE_ATTRIBUTES); } - private void assertOutputContains(String... messages) { + private void assertErrorContains(String... messages) { for (String message : messages) { if (!errOutput.contains(message)) { - fail("Expected output to contain '" + message - + "' but err_output was:\n" + errOutput); + fail( + "Expected output to contain '" + message + "' but err_output was:\n" + + errOutput); + } + } + } + + private void assertSysOutContains(String... messages) { + for (String message : messages) { + if (!sysOutput.contains(message)) { + fail( + "Expected output to contain '" + message + "' but sys_output was:\n" + + sysOutput); } } } private int runTool(String... args) throws Exception { errOutBytes.reset(); + sysOutBytes.reset(); LOG.info("Running: NodeAttributesCLI " + Joiner.on(" ").join(args)); int ret = nodeAttributesCLI.run(args); errOutput = new String(errOutBytes.toByteArray(), Charsets.UTF_8); + sysOutput = new String(sysOutBytes.toByteArray(), Charsets.UTF_8); LOG.info("Err_output:\n" + errOutput); + LOG.info("Sys_output:\n" + sysOutput); return ret; } } http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java index 6b823b2..3623908 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java @@ -17,6 +17,8 @@ */ package org.apache.hadoop.yarn.client.cli; +import org.apache.hadoop.yarn.api.records.NodeAttribute; +import org.apache.hadoop.yarn.api.records.NodeAttributeType; import static org.junit.Assert.assertEquals; import static org.mockito.Matchers.any; import static org.mockito.Matchers.anyInt; @@ -1543,8 +1545,8 @@ public class TestYarnCLI { public void testNodeStatus() throws Exception { NodeId nodeId = NodeId.newInstance("host0", 0); NodeCLI cli = new NodeCLI(); - when(client.getNodeReports()).thenReturn( - getNodeReports(3, NodeState.RUNNING, false)); + when(client.getNodeReports()) + .thenReturn(getNodeReports(3, NodeState.RUNNING, false, false, false)); cli.setClient(client); cli.setSysOutPrintStream(sysOut); cli.setSysErrPrintStream(sysErr); @@ -1567,6 +1569,8 @@ public class TestYarnCLI { pw.println("\tCPU-Used : 0 vcores"); pw.println("\tCPU-Capacity : 0 vcores"); pw.println("\tNode-Labels : a,b,c,x,y,z"); + pw.println("\tNode Attributes : rm.yarn.io/GPU(STRING)=ARM"); + pw.println("\t rm.yarn.io/CPU(STRING)=ARM"); pw.println("\tResource Utilization by Node : PMem:2048 MB, VMem:4096 MB, VCores:8.0"); pw.println("\tResource Utilization by Containers : PMem:1024 MB, VMem:2048 MB, VCores:4.0"); pw.close(); @@ -1603,6 +1607,7 @@ public class TestYarnCLI { pw.println("\tCPU-Used : 0 vcores"); pw.println("\tCPU-Capacity : 0 vcores"); pw.println("\tNode-Labels : "); + pw.println("\tNode Attributes : "); pw.println("\tResource Utilization by Node : PMem:2048 MB, VMem:4096 MB, VCores:8.0"); pw.println("\tResource Utilization by Containers : PMem:1024 MB, VMem:2048 MB, VCores:4.0"); pw.close(); @@ -1615,8 +1620,8 @@ public class TestYarnCLI { public void testNodeStatusWithEmptyResourceUtilization() throws Exception { NodeId nodeId = NodeId.newInstance("host0", 0); NodeCLI cli = new NodeCLI(); - when(client.getNodeReports()).thenReturn( - getNodeReports(3, NodeState.RUNNING, false, true)); + when(client.getNodeReports()) + .thenReturn(getNodeReports(3, NodeState.RUNNING, false, true, true)); cli.setClient(client); cli.setSysOutPrintStream(sysOut); cli.setSysErrPrintStream(sysErr); @@ -1639,6 +1644,7 @@ public class TestYarnCLI { pw.println("\tCPU-Used : 0 vcores"); pw.println("\tCPU-Capacity : 0 vcores"); pw.println("\tNode-Labels : a,b,c,x,y,z"); + pw.println("\tNode Attributes : "); pw.println("\tResource Utilization by Node : "); pw.println("\tResource Utilization by Containers : "); pw.close(); @@ -2048,18 +2054,20 @@ public class TestYarnCLI { cli.run(new String[] { "application" }); verify(sysErr).println("Invalid Command Usage : "); } - + private List<NodeReport> getNodeReports(int noOfNodes, NodeState state) { - return getNodeReports(noOfNodes, state, true, false); + return getNodeReports(noOfNodes, state, true, false, true); } private List<NodeReport> getNodeReports(int noOfNodes, NodeState state, - boolean emptyNodeLabel) { - return getNodeReports(noOfNodes, state, emptyNodeLabel, false); + boolean emptyNodeLabel, boolean emptyAttributes) { + return getNodeReports(noOfNodes, state, emptyNodeLabel, false, + emptyAttributes); } private List<NodeReport> getNodeReports(int noOfNodes, NodeState state, - boolean emptyNodeLabel, boolean emptyResourceUtilization) { + boolean emptyNodeLabel, boolean emptyResourceUtilization, + boolean emptyAttributes) { List<NodeReport> nodeReports = new ArrayList<NodeReport>(); for (int i = 0; i < noOfNodes; i++) { @@ -2081,6 +2089,11 @@ public class TestYarnCLI { nodeReport.setAggregatedContainersUtilization(containersUtilization); nodeReport.setNodeUtilization(nodeUtilization); } + if (!emptyAttributes) { + nodeReport.setNodeAttributes(ImmutableSet.of(NodeAttribute + .newInstance("GPU", NodeAttributeType.STRING, "ARM"), + NodeAttribute.newInstance("CPU", NodeAttributeType.STRING, "ARM"))); + } nodeReports.add(nodeReport); } return nodeReports; http://git-wip-us.apache.org/repos/asf/hadoop/blob/acd7729a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributeInfoPBImpl.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributeInfoPBImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributeInfoPBImpl.java index bff6335..e2db568 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributeInfoPBImpl.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributeInfoPBImpl.java @@ -130,14 +130,18 @@ public class NodeAttributeInfoPBImpl extends NodeAttributeInfo { } if (obj instanceof NodeAttributeInfo) { NodeAttributeInfo other = (NodeAttributeInfo) obj; - getAttributeKey().equals(other.getAttributeKey()); - return true; + return getAttributeKey().equals(other.getAttributeKey()); } return false; } @Override public String toString() { - return getAttributeKey().toString() + ":Type-" + getAttributeType(); + StringBuilder strBuilder = new StringBuilder(); + NodeAttributeKey key = this.getAttributeKey(); + strBuilder.append(key.getAttributePrefix()).append("/") + .append(key.getAttributeName()).append("(") + .append(this.getAttributeType()).append(")"); + return strBuilder.toString(); } } --------------------------------------------------------------------- To unsubscribe, e-mail: common-commits-unsubscr...@hadoop.apache.org For additional commands, e-mail: common-commits-h...@hadoop.apache.org
