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/Cluster index b930014..9fe657c 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 @@ -25,17 +25,23 @@ import java.io.UnsupportedEncodingException; import java.nio.charset.Charset; import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; 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,6 +54,8 @@ */ @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 = @@ -55,6 +63,7 @@ public static final String CMD = "cluster"; private boolean accessLocal = false; static CommonNodeLabelsManager localNodeLabelsManager = null; + public static final String NODES = "nodes"; public static void main(String[] args) throws Exception { ClusterCLI cli = new ClusterCLI(); @@ -70,7 +79,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 +95,7 @@ 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 { @@ -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/Tes index 5a0f049..6800df8 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; @@ -103,7 +107,66 @@ public void testGetClusterNodeLabelsWithLocalAccess() 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 testGetEmptyClusterNodeLabels() throws Exception { YarnClient client = mock(YarnClient.class); @@ -133,7 +196,8 @@ public void testHelp() throws Exception { int rc = cli.run(new String[] { "cluster", "--help" }); assertEquals(0, rc); - + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); PrintWriter pw = new PrintWriter(baos); pw.println("usage: yarn cluster"); @@ -158,7 +222,12 @@ 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")); }