YARN-8103. Add CLI interface to query node attributes. Contributed by Bibin A Chundatt.

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

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();
   }
 

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'
 	;;

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();
 }

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 {

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) {

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);
   }
 }

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:"

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();

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;
   }
 }

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;
@@ -1544,8 +1546,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);
@@ -1568,6 +1570,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();
@@ -1604,6 +1608,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();
@@ -1616,8 +1621,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);
@@ -1640,6 +1645,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();
@@ -2049,18 +2055,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++) {
@@ -2082,6 +2090,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;

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();
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributePBImpl.java

@@ -152,15 +152,19 @@ public class NodeAttributePBImpl extends NodeAttribute {
     }
     if (obj instanceof NodeAttribute) {
       NodeAttribute other = (NodeAttribute) obj;
-      getAttributeKey().equals(other.getAttributeKey());
-      return true;
+      return getAttributeKey().equals(other.getAttributeKey());
     }
     return false;
   }
 
   @Override
   public String toString() {
-    return getAttributeKey().toString() + ":Value-" + getAttributeValue()
-        + ":Type-" + getAttributeType();
+    StringBuilder strBuilder = new StringBuilder();
+    NodeAttributeKey key = this.getAttributeKey();
+    strBuilder.append(key.getAttributePrefix()).append("/")
+        .append(key.getAttributeName()).append("(")
+        .append(this.getAttributeType()).append(")=")
+        .append(this.getAttributeValue());
+    return strBuilder.toString();
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeReportPBImpl.java

@@ -18,17 +18,21 @@
 
 package org.apache.hadoop.yarn.api.records.impl.pb;
 
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 import org.apache.hadoop.classification.InterfaceAudience.Private;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
+import org.apache.hadoop.yarn.api.records.NodeAttribute;
 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.api.records.NodeUpdateType;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceUtilization;
+import org.apache.hadoop.yarn.proto.YarnProtos.NodeAttributeProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.NodeIdProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.NodeReportProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.NodeReportProtoOrBuilder;
@@ -50,6 +54,7 @@ public class NodeReportPBImpl extends NodeReport {
   private ResourceUtilization containersUtilization = null;
   private ResourceUtilization nodeUtilization = null;
   Set<String> labels;
+  private Set<NodeAttribute> nodeAttributes;
 
   public NodeReportPBImpl() {
     builder = NodeReportProto.newBuilder();
@@ -268,6 +273,14 @@ public class NodeReportPBImpl extends NodeReport {
       builder.clearNodeLabels();
       builder.addAllNodeLabels(this.labels);
     }
+    if (this.nodeAttributes != null) {
+      builder.clearNodeAttributes();
+      List<NodeAttributeProto> attrList = new ArrayList<>();
+      for (NodeAttribute attr : this.nodeAttributes) {
+        attrList.add(convertToProtoFormat(attr));
+      }
+      builder.addAllNodeAttributes(attrList);
+    }
     if (this.nodeUtilization != null
         && !((ResourceUtilizationPBImpl) this.nodeUtilization).getProto()
             .equals(builder.getNodeUtilization())) {
@@ -306,7 +319,16 @@ public class NodeReportPBImpl extends NodeReport {
   private NodeIdProto convertToProtoFormat(NodeId nodeId) {
     return ((NodeIdPBImpl) nodeId).getProto();
   }
-  
+
+  private NodeAttributeProto convertToProtoFormat(NodeAttribute nodeAttr) {
+    return ((NodeAttributePBImpl) nodeAttr).getProto();
+  }
+
+  private NodeAttributePBImpl convertFromProtoFormat(
+      NodeAttributeProto nodeAttr) {
+    return new NodeAttributePBImpl(nodeAttr);
+  }
+
   private ResourcePBImpl convertFromProtoFormat(ResourceProto p) {
     return new ResourcePBImpl(p);
   }
@@ -427,4 +449,24 @@ public class NodeReportPBImpl extends NodeReport {
     }
     builder.setNodeUpdateType(ProtoUtils.convertToProtoFormat(nodeUpdateType));
   }
+
+  @Override
+  public void setNodeAttributes(Set<NodeAttribute> nodeAttrs) {
+    maybeInitBuilder();
+    builder.clearNodeAttributes();
+    this.nodeAttributes = nodeAttrs;
+  }
+
+  @Override
+  public Set<NodeAttribute> getNodeAttributes() {
+    if (nodeAttributes != null) {
+      return nodeAttributes;
+    }
+    NodeReportProtoOrBuilder p = viaProto ? proto : builder;
+    this.nodeAttributes = new HashSet<>();
+    for (NodeAttributeProto nattrProto : p.getNodeAttributesList()) {
+      nodeAttributes.add(convertFromProtoFormat(nattrProto));
+    }
+    return nodeAttributes;
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/utils/BuilderUtils.java

@@ -54,6 +54,7 @@ import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.api.records.LocalResourceType;
 import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
+import org.apache.hadoop.yarn.api.records.NodeAttribute;
 import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.api.records.NodeReport;
 import org.apache.hadoop.yarn.api.records.NodeState;
@@ -199,7 +200,7 @@ public class BuilderUtils {
       NodeUpdateType nodeUpdateType) {
     return newNodeReport(nodeId, nodeState, httpAddress, rackName, used,
         capability, numContainers, healthReport, lastHealthReportTime,
-        nodeLabels, null, null, decommissioningTimeout, nodeUpdateType);
+        nodeLabels, null, null, decommissioningTimeout, nodeUpdateType, null);
   }
 
   public static NodeReport newNodeReport(NodeId nodeId, NodeState nodeState,
@@ -207,7 +208,7 @@ public class BuilderUtils {
       int numContainers, String healthReport, long lastHealthReportTime,
       Set<String> nodeLabels, ResourceUtilization containersUtilization,
       ResourceUtilization nodeUtilization, Integer decommissioningTimeout,
-      NodeUpdateType nodeUpdateType) {
+      NodeUpdateType nodeUpdateType, Set<NodeAttribute> attrs) {
     NodeReport nodeReport = recordFactory.newRecordInstance(NodeReport.class);
     nodeReport.setNodeId(nodeId);
     nodeReport.setNodeState(nodeState);
@@ -223,6 +224,7 @@ public class BuilderUtils {
     nodeReport.setNodeUtilization(nodeUtilization);
     nodeReport.setDecommissioningTimeout(decommissioningTimeout);
     nodeReport.setNodeUpdateType(nodeUpdateType);
+    nodeReport.setNodeAttributes(attrs);
     return nodeReport;
   }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/AdminService.java

@@ -983,12 +983,11 @@ public class AdminService extends CompositeService implements
     List<NodeToAttributes> nodesToAttributes = request.getNodesToAttributes();
     boolean failOnUnknownNodes = request.getFailOnUnknownNodes();
 
-    Map<String, Set<NodeAttribute>> nodeAttributeMapping =
-        validateAndFetch(nodesToAttributes, failOnUnknownNodes);
-
     NodeAttributesManager nodeAttributesManager =
         rm.getRMContext().getNodeAttributesManager();
     try {
+      Map<String, Set<NodeAttribute>> nodeAttributeMapping =
+          validateAndFetch(nodesToAttributes, failOnUnknownNodes);
       switch (request.getOperation()) {
       case ADD:
         nodeAttributesManager.addNodeAttributes(nodeAttributeMapping);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/ClientRMService.java

@@ -1050,8 +1050,9 @@ public class ClientRMService extends AbstractService implements
     if (schedulerNodeReport != null) {
       used = schedulerNodeReport.getUsedResource();
       numContainers = schedulerNodeReport.getNumContainers();
-    } 
+    }
 
+    Set<NodeAttribute> attrs = rmNode.getAllNodeAttributes();
     NodeReport report =
         BuilderUtils.newNodeReport(rmNode.getNodeID(), rmNode.getState(),
             rmNode.getHttpAddress(), rmNode.getRackName(), used,
@@ -1059,7 +1060,7 @@ public class ClientRMService extends AbstractService implements
             rmNode.getHealthReport(), rmNode.getLastHealthReportTime(),
             rmNode.getNodeLabels(), rmNode.getAggregatedContainersUtilization(),
             rmNode.getNodeUtilization(), rmNode.getDecommissioningTimeout(),
-            null);
+            null, attrs);
 
     return report;
   }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/ResourceTrackerService.java

@@ -673,10 +673,6 @@ public class ResourceTrackerService extends AbstractService implements
         this.rmContext.getNodeAttributesManager()
             .replaceNodeAttributes(NodeAttribute.PREFIX_DISTRIBUTED,
                 ImmutableMap.of(nodeId.getHost(), nodeAttributes));
-
-        // Update node attributes to RMNode
-        rmNode.setNodeAttributes(NodeAttribute.PREFIX_DISTRIBUTED,
-            nodeAttributes);
       }
     }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNode.java

@@ -197,15 +197,8 @@ public interface RMNode {
    */
   RMContext getRMContext();
 
- /**
-   * Sets node attributes per prefix.
-   * @param prefix node attribute prefix
-   * @param nodeAttributes node attributes
-   */
-  void setNodeAttributes(String prefix, Set<NodeAttribute> nodeAttributes);
-
   /**
-   * @return all node attributes grouped by their prefix as a map.
+   * @return all node attributes as a Set.
    */
-  Map<String, Set<NodeAttribute>> getAllNodeAttributes();
+  Set<NodeAttribute> getAllNodeAttributes();
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java

@@ -59,7 +59,9 @@ import org.apache.hadoop.yarn.api.records.ResourceUtilization;
 import org.apache.hadoop.yarn.event.EventHandler;
 import org.apache.hadoop.yarn.factories.RecordFactory;
 import org.apache.hadoop.yarn.factory.providers.RecordFactoryProvider;
+import org.apache.hadoop.yarn.nodelabels.AttributeValue;
 import org.apache.hadoop.yarn.nodelabels.CommonNodeLabelsManager;
+import org.apache.hadoop.yarn.nodelabels.NodeAttributesManager;
 import org.apache.hadoop.yarn.server.api.protocolrecords.LogAggregationReport;
 import org.apache.hadoop.yarn.server.api.protocolrecords.NMContainerStatus;
 import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatResponse;
@@ -186,9 +188,6 @@ public class RMNodeImpl implements RMNode, EventHandler<RMNodeEvent> {
   private NodeHeartbeatResponse latestNodeHeartBeatResponse = recordFactory
       .newRecordInstance(NodeHeartbeatResponse.class);
 
-  // Node attributes, store by prefix
-  private Map<String, Set<NodeAttribute>> nodeAttributes = new HashMap<>();
-
   private static final StateMachineFactory<RMNodeImpl,
                                            NodeState,
                                            RMNodeEventType,
@@ -1552,13 +1551,10 @@ public class RMNodeImpl implements RMNode, EventHandler<RMNodeEvent> {
   }
 
   @Override
-  public void setNodeAttributes(String prefix,
-      Set<NodeAttribute> nodeAttributeSet) {
-    this.nodeAttributes.put(prefix, nodeAttributeSet);
-  }
-
-  @Override
-  public Map<String, Set<NodeAttribute>> getAllNodeAttributes() {
-    return this.nodeAttributes;
+  public Set<NodeAttribute> getAllNodeAttributes() {
+    NodeAttributesManager attrMgr = context.getNodeAttributesManager();
+    Map<NodeAttribute, AttributeValue> nodeattrs =
+        attrMgr.getAttributesForNode(hostName);
+    return nodeattrs.keySet();
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/dao/NodeInfo.java

@@ -116,16 +116,11 @@ public class NodeInfo {
     }
 
     // add attributes
-    Map<String, Set<NodeAttribute>> nodeAttributes =
-        ni.getAllNodeAttributes();
+    Set<NodeAttribute> attrs = ni.getAllNodeAttributes();
     nodeAttributesInfo = new NodeAttributesInfo();
-    if (nodeAttributes != null) {
-      for (Set<NodeAttribute> attrs : nodeAttributes.values()) {
-        for (NodeAttribute attribute : attrs) {
-          NodeAttributeInfo info = new NodeAttributeInfo(attribute);
-          this.nodeAttributesInfo.addNodeAttributeInfo(info);
-        }
-      }
+    for (NodeAttribute attribute : attrs) {
+      NodeAttributeInfo info = new NodeAttributeInfo(attribute);
+      this.nodeAttributesInfo.addNodeAttributeInfo(info);
     }
 
     // add allocation tags

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNodes.java

@@ -292,8 +292,8 @@ public class MockNodes {
     }
 
     @Override
-    public Map<String, Set<NodeAttribute>> getAllNodeAttributes() {
-      return null;
+    public Set<NodeAttribute> getAllNodeAttributes() {
+      return Collections.emptySet();
     }
 
     @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMAdminService.java

@@ -1611,9 +1611,10 @@ public class TestRMAdminService {
     try {
       rm.adminService.mapAttributesToNodes(request);
       fail("host5 is not a valid node, It should have failed");
-    } catch (Exception ex) {
+    } catch (YarnException ex) {
       Assert.assertEquals("Exception Message is not as desired",
-          " Following nodes does not exist : [host5]", ex.getMessage());
+          " Following nodes does not exist : [host5]",
+          ex.getCause().getMessage());
     }
 
     request =
@@ -1633,10 +1634,10 @@ public class TestRMAdminService {
       // against hostname hence the message as : nodes does not exist.
       rm.adminService.mapAttributesToNodes(request);
       fail("host with the port should fail as only hostnames are validated");
-    } catch (Exception ex) {
+    } catch (YarnException ex) {
       Assert.assertEquals("Exception Message is not as desired",
           " Following nodes does not exist : [host4:8889, host2:8889]",
-          ex.getMessage());
+          ex.getCause().getMessage());
     }
 
     request =
@@ -1669,11 +1670,10 @@ public class TestRMAdminService {
     try {
       rm.adminService.mapAttributesToNodes(request);
       fail("This operation should fail as prefix should be \"nm.yarn.io\".");
-    } catch (Exception ex) {
+    } catch (YarnException ex) {
       Assert.assertEquals("Exception Message is not as desired",
           "Invalid Attribute Mapping for the node host5. Prefix should be "
-              + "rm.yarn.io",
-          ex.getMessage());
+              + "rm.yarn.io", ex.getCause().getMessage());
     }
 
     rm.close();