diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java index 7ead774..745a98b 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ClusterCLI.java @@ -26,6 +26,8 @@ import java.nio.charset.Charset; import java.util.ArrayList; import java.util.List; +import java.util.Map; +import java.util.Set; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.GnuParser; @@ -37,6 +39,7 @@ 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.NodeId; import org.apache.hadoop.yarn.api.records.NodeLabel; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; @@ -49,12 +52,15 @@ */ @Private public class ClusterCLI extends YarnCLI { + public static final String LABELS_PATTERN = "%32s\t%s" + + System.getProperty("line.separator"); private static final String TITLE = "yarn cluster"; 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"; + public static final String NODES = "nodes"; private boolean accessLocal = false; static CommonNodeLabelsManager localNodeLabelsManager = null; @@ -72,7 +78,8 @@ public int run(String[] args) throws Exception { Options opts = new Options(); opts.addOption("lnl", LIST_LABELS_CMD, false, - "List cluster node-label collection"); + "List cluster node-label collection. " + + "Supports optional use of -nodes to print all nodes partitioned by labels."); opts.addOption("lna", LIST_CLUSTER_ATTRIBUTES, false, "List cluster node-attribute collection"); opts.addOption("h", HELP_CMD, false, "Displays help for all commands."); @@ -89,7 +96,8 @@ public int run(String[] args) throws Exception { + " when the command run on the machine where RM is running." + " Also, this option is UNSTABLE, could be removed in future" + " releases."); - + opts.addOption(NODES, false, "Work with -lnl/list-node-labels" + + " to print all nodes partitioned by labels."); int exitCode = -1; CommandLine parsedCli = null; try { @@ -107,7 +115,11 @@ public int run(String[] args) throws Exception { } if (parsedCli.hasOption(LIST_LABELS_CMD)) { - printClusterNodeLabels(); + if (parsedCli.hasOption(NODES)) { + printClusterNodeLabelsMap(); + } else { + printClusterNodeLabels(); + } } else if(parsedCli.hasOption(LIST_CLUSTER_ATTRIBUTES)){ printClusterNodeAttributes(); } else if (parsedCli.hasOption(HELP_CMD)) { @@ -131,6 +143,26 @@ private void printClusterNodeAttributes() throws IOException, YarnException { sysout.println(baos.toString("UTF-8")); } + void printClusterNodeLabelsMap()throws YarnException, IOException { + Map> labelsMap= null; + if (accessLocal) { + labelsMap = getNodeLabelManagerInstance(getConf()).getLabelsToNodes(); + } else { + labelsMap = client.getLabelsToNodes(); + } + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter pw = + new PrintWriter(new OutputStreamWriter(baos, Charset.forName("UTF-8"))); + pw.println("Node Labels Num: "+ labelsMap.size()); + pw.printf(LABELS_PATTERN,"Labels", "Nodes"); + for(String label : labelsMap.keySet()) { + pw.printf(LABELS_PATTERN, NodeLabel.newInstance(label).toString(), + StringUtils.join(labelsMap.get(label).iterator(), ", ")); + } + pw.close(); + sysout.println(baos.toString("UTF-8")); + } + void printClusterNodeLabels() throws YarnException, IOException { List nodeLabels = null; if (accessLocal) { diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java index 85aca05..156a615 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java @@ -18,6 +18,7 @@ package org.apache.hadoop.yarn.client.cli; +import com.google.common.collect.Sets; import org.apache.hadoop.yarn.api.records.NodeAttributeInfo; import org.apache.hadoop.yarn.api.records.NodeAttributeKey; import org.apache.hadoop.yarn.api.records.NodeAttributeType; @@ -32,7 +33,12 @@ import java.io.PrintWriter; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import org.apache.hadoop.yarn.api.records.NodeId; import org.apache.hadoop.yarn.api.records.NodeLabel; import org.apache.hadoop.yarn.client.api.YarnClient; import org.apache.hadoop.yarn.nodelabels.CommonNodeLabelsManager; @@ -137,7 +143,67 @@ public void testGetEmptyClusterNodeLabels() throws Exception { pw.close(); verify(sysOut).println(baos.toString("UTF-8")); } - + + @Test + public void testGetNodesCollectionByLabel() throws Exception { + YarnClient client = mock(YarnClient.class); + Map> nodes = new HashMap<>(); + Set nodeset = new HashSet(); + nodeset.add(NodeId.newInstance("node1", 45454)); + nodeset.add(NodeId.newInstance("node2", 45454)); + nodes.put("label1", nodeset); + nodes.put("label2", Sets.newHashSet(NodeId.newInstance("node3", 45454))); + when(client.getLabelsToNodes()).thenReturn(nodes); + ClusterCLI cli = new ClusterCLI(); + cli.setClient(client); + cli.setSysOutPrintStream(sysOut); + cli.setSysErrPrintStream(sysErr); + + int rc = + cli.run(new String[] { ClusterCLI.CMD, "-" + ClusterCLI.LIST_LABELS_CMD ,"-" + ClusterCLI.NODES}); + assertEquals(0, rc); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter pw = new PrintWriter(baos); + pw.println("Node Labels Num: 2"); + pw.printf(ClusterCLI.LABELS_PATTERN, "Labels", "Nodes"); + pw.printf(ClusterCLI.LABELS_PATTERN, "", "node2:45454, node1:45454"); + pw.printf(ClusterCLI.LABELS_PATTERN, "", "node3:45454"); + pw.close(); + verify(sysOut).print(baos.toString("UTF-8")); + } + + @Test + public void testGetNodesCollectionByLabelWithLocalAccess() throws Exception { + YarnClient client = mock(YarnClient.class); + Map> nodes = new HashMap<>(); + Set nodeset = new HashSet(); + nodeset.add(NodeId.newInstance("node1", 45454)); + nodeset.add(NodeId.newInstance("node2", 45454)); + nodes.put("label1", nodeset); + nodes.put("label2", Sets.newHashSet(NodeId.newInstance("node3", 45454))); + when(client.getLabelsToNodes()).thenReturn(nodes); + ClusterCLI cli = new ClusterCLI(); + cli.setClient(client); + cli.setSysOutPrintStream(sysOut); + cli.setSysErrPrintStream(sysErr); + ClusterCLI.localNodeLabelsManager = mock(CommonNodeLabelsManager.class); + when(ClusterCLI.localNodeLabelsManager.getLabelsToNodes()).thenReturn(nodes); + + int rc = + cli.run(new String[] { ClusterCLI.CMD, "-" + ClusterCLI.LIST_LABELS_CMD ,"-" + ClusterCLI.NODES, + "-" + ClusterCLI.DIRECTLY_ACCESS_NODE_LABEL_STORE}); + assertEquals(0, rc); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + PrintWriter pw = new PrintWriter(baos); + pw.println("Node Labels Num: 2"); + pw.printf(ClusterCLI.LABELS_PATTERN, "Labels", "Nodes"); + pw.printf(ClusterCLI.LABELS_PATTERN, "", "node2:45454, node1:45454"); + pw.printf(ClusterCLI.LABELS_PATTERN, "", "node3:45454"); + pw.close(); + verify(sysOut).print(baos.toString("UTF-8")); + } + + @Test public void testHelp() throws Exception { ClusterCLI cli = createAndGetClusterCLI(); @@ -172,7 +238,12 @@ public void testHelp() throws Exception { 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.println(" collection. Supports optional"); + pw.println(" use of -nodes to print all"); + pw.println(" nodes partitioned by labels."); + pw.println(" -nodes Work with -lnl/list-node-labels"); + pw.println(" to print all nodes partitioned"); + pw.println(" by labels."); pw.close(); verify(sysOut).println(baos.toString("UTF-8")); }