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 b930014..8b15753 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,16 +26,20 @@ 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; import org.apache.commons.cli.HelpFormatter; import org.apache.commons.cli.MissingArgumentException; +import org.apache.commons.cli.Option; import org.apache.commons.cli.Options; import org.apache.commons.lang.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.NodeId; import org.apache.hadoop.yarn.api.records.NodeLabel; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; @@ -48,10 +52,13 @@ */ @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 NODES = "nodes"; public static final String CMD = "cluster"; private boolean accessLocal = false; static CommonNodeLabelsManager localNodeLabelsManager = null; @@ -70,7 +77,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("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, " @@ -85,7 +93,9 @@ 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."); - + Option nodesOpt = new Option(NODES, false, + "Work with -lnl/list-node-labels to print all nodes partitioned by labels."); + opts.addOption(nodesOpt); int exitCode = -1; CommandLine parsedCli = null; try { @@ -101,7 +111,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(HELP_CMD)) { printUsage(opts); return 0; @@ -124,6 +138,26 @@ void printClusterNodeLabels() throws YarnException, IOException { StringUtils.join(nodeLabels.iterator(), ","))); } + 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")); + } + @VisibleForTesting static synchronized CommonNodeLabelsManager getNodeLabelManagerInstance(Configuration conf) { diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestClusterCLI.java index 5a0f049..758e6b3 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 @@ -29,15 +29,19 @@ 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 com.google.common.collect.Sets; +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; import org.junit.Before; import org.junit.Test; -import com.google.common.collect.ImmutableSet; - public class TestClusterCLI { ByteArrayOutputStream sysOutStream; private PrintStream sysOut; @@ -74,7 +78,67 @@ public void testGetClusterNodeLabels() 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 testGetClusterNodeLabelsWithLocalAccess() throws Exception { YarnClient client = mock(YarnClient.class); @@ -158,7 +222,13 @@ public void testHelp() throws Exception { pw.println(" removed in future releases."); pw.println(" -h,--help Displays help for all commands."); 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")); }