YARN-6856. [YARN-3409] Support CLI for Node Attributes Mapping. Contributed by Naganarasimha G R.
This commit is contained in:
Naganarasimha
Sunil G
parent
1f42ce907a
commit
2475fb0a1e
@ -575,7 +575,7 @@ private int help(String[] argv) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
protected static class UsageInfo {
|
||||
public static class UsageInfo {
|
||||
public final String args;
|
||||
public final String help;
|
||||
|
||||
|
||||
@ -55,6 +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 "version" client "print the version"
|
||||
hadoop_generate_usage "${HADOOP_SHELL_EXECNAME}" true
|
||||
}
|
||||
@ -186,6 +187,10 @@ ${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)
|
||||
HADOOP_SUBCMD_SUPPORTDAEMONIZATION="false"
|
||||
HADOOP_CLASSNAME='org.apache.hadoop.yarn.client.cli.NodeAttributesCLI'
|
||||
;;
|
||||
timelineserver)
|
||||
HADOOP_SUBCMD_SUPPORTDAEMONIZATION="true"
|
||||
HADOOP_CLASSNAME='org.apache.hadoop.yarn.server.applicationhistoryservice.ApplicationHistoryServer'
|
||||
|
||||
@ -0,0 +1,410 @@
|
||||
/**
|
||||
* Licensed to the Apache Software Foundation (ASF) under one
|
||||
* or more contributor license agreements. See the NOTICE file
|
||||
* distributed with this work for additional information
|
||||
* regarding copyright ownership. The ASF licenses this file
|
||||
* to you under the Apache License, Version 2.0 (the
|
||||
* "License"); you may not use this file except in compliance
|
||||
* with the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.apache.hadoop.yarn.client.cli;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.PrintStream;
|
||||
import java.util.ArrayList;
|
||||
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.MissingArgumentException;
|
||||
import org.apache.commons.cli.Options;
|
||||
import org.apache.commons.cli.ParseException;
|
||||
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.records.NodeAttribute;
|
||||
import org.apache.hadoop.yarn.api.records.NodeAttributeType;
|
||||
import org.apache.hadoop.yarn.client.ClientRMProxy;
|
||||
import org.apache.hadoop.yarn.conf.YarnConfiguration;
|
||||
import org.apache.hadoop.yarn.exceptions.YarnException;
|
||||
import org.apache.hadoop.yarn.server.api.ResourceManagerAdministrationProtocol;
|
||||
import org.apache.hadoop.yarn.server.api.protocolrecords.AttributeMappingOperationType;
|
||||
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;
|
||||
|
||||
/**
|
||||
* CLI to map attributes to Nodes.
|
||||
*
|
||||
*/
|
||||
public class NodeAttributesCLI extends Configured implements Tool {
|
||||
|
||||
protected static final String INVALID_MAPPING_ERR_MSG =
|
||||
"Invalid Node to attribute mapping : ";
|
||||
|
||||
protected static final String USAGE_YARN_NODE_ATTRIBUTES =
|
||||
"Usage: yarn node-attributes ";
|
||||
|
||||
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 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);
|
||||
}
|
||||
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
||||
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");
|
||||
}
|
||||
builder.append(" -help" + " [cmd]\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Displays format of commands.
|
||||
*
|
||||
* @param cmd The command that is being executed.
|
||||
*/
|
||||
private void printUsage(String cmd) {
|
||||
StringBuilder usageBuilder = new StringBuilder();
|
||||
if (NODE_ATTRIB_USAGE.containsKey(cmd)) {
|
||||
buildIndividualUsageMsg(cmd, usageBuilder);
|
||||
} else {
|
||||
buildUsageMsgForAllCmds(usageBuilder);
|
||||
}
|
||||
errOut.println(usageBuilder);
|
||||
}
|
||||
|
||||
private void printUsage() {
|
||||
printUsage("");
|
||||
}
|
||||
|
||||
protected ResourceManagerAdministrationProtocol createAdminProtocol()
|
||||
throws IOException {
|
||||
// Get the current configuration
|
||||
final YarnConfiguration conf = new YarnConfiguration(getConf());
|
||||
return ClientRMProxy.createRMProxy(conf,
|
||||
ResourceManagerAdministrationProtocol.class);
|
||||
}
|
||||
|
||||
@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 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;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int run(String[] args) throws Exception {
|
||||
if (args.length < 1) {
|
||||
printUsage();
|
||||
return -1;
|
||||
}
|
||||
|
||||
int exitCode = -1;
|
||||
int i = 0;
|
||||
String cmd = args[i++];
|
||||
|
||||
if ("-help".equals(cmd)) {
|
||||
exitCode = 0;
|
||||
if (args.length >= 2) {
|
||||
printHelpMsg(args[i]);
|
||||
} else {
|
||||
printHelpMsg("");
|
||||
}
|
||||
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);
|
||||
} else {
|
||||
exitCode = -1;
|
||||
errOut.println(cmd.substring(1) + ": Unknown command");
|
||||
printUsage();
|
||||
}
|
||||
} 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());
|
||||
}
|
||||
} catch (Exception e) {
|
||||
exitCode = -1;
|
||||
errOut.println(cmd.substring(1) + ": " + e.getLocalizedMessage());
|
||||
}
|
||||
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;
|
||||
}
|
||||
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
|
||||
*/
|
||||
private List<NodeToAttributes> buildNodeLabelsMapFromStr(String args,
|
||||
boolean validateForAttributes, AttributeMappingOperationType operation) {
|
||||
List<NodeToAttributes> nodeToAttributesList = new ArrayList<>();
|
||||
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);
|
||||
} 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);
|
||||
}
|
||||
}
|
||||
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(attributeName.trim(),
|
||||
attributeType, attributeValue.trim()));
|
||||
}
|
||||
}
|
||||
if (validateForAttributes) {
|
||||
Preconditions.checkArgument((attributesList.size() > 0),
|
||||
"Attributes cannot be null or empty for Operation "
|
||||
+ operation.name() + " on the node " + node);
|
||||
}
|
||||
nodeToAttributesList
|
||||
.add(NodeToAttributes.newInstance(node, attributesList));
|
||||
}
|
||||
|
||||
if (nodeToAttributesList.isEmpty()) {
|
||||
throw new IllegalArgumentException(NO_MAPPING_ERR_MSG);
|
||||
}
|
||||
return nodeToAttributesList;
|
||||
}
|
||||
|
||||
public static void main(String[] args) throws Exception {
|
||||
int result = ToolRunner.run(new NodeAttributesCLI(), args);
|
||||
System.exit(result);
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,328 @@
|
||||
/**
|
||||
* Licensed to the Apache Software Foundation (ASF) under one
|
||||
* or more contributor license agreements. See the NOTICE file
|
||||
* distributed with this work for additional information
|
||||
* regarding copyright ownership. The ASF licenses this file
|
||||
* to you under the Apache License, Version 2.0 (the
|
||||
* "License"); you may not use this file except in compliance
|
||||
* with the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.apache.hadoop.yarn.client.cli;
|
||||
|
||||
import static org.junit.Assert.assertTrue;
|
||||
import static org.junit.Assert.fail;
|
||||
import static org.mockito.Matchers.any;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.PrintStream;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
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;
|
||||
import org.apache.hadoop.yarn.server.api.ResourceManagerAdministrationProtocol;
|
||||
import org.apache.hadoop.yarn.server.api.protocolrecords.AttributeMappingOperationType;
|
||||
import org.apache.hadoop.yarn.server.api.protocolrecords.NodeToAttributes;
|
||||
import org.apache.hadoop.yarn.server.api.protocolrecords.NodesToAttributesMappingRequest;
|
||||
import org.apache.hadoop.yarn.server.api.protocolrecords.NodesToAttributesMappingResponse;
|
||||
import org.junit.Before;
|
||||
import org.junit.Test;
|
||||
import org.mockito.invocation.InvocationOnMock;
|
||||
import org.mockito.stubbing.Answer;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import com.google.common.base.Charsets;
|
||||
import com.google.common.base.Joiner;
|
||||
|
||||
/**
|
||||
* Test class for TestNodeAttributesCLI.
|
||||
*/
|
||||
public class TestNodeAttributesCLI {
|
||||
private static final Logger LOG =
|
||||
LoggerFactory.getLogger(TestNodeAttributesCLI.class);
|
||||
private ResourceManagerAdministrationProtocol admin;
|
||||
private NodesToAttributesMappingRequest request;
|
||||
private NodeAttributesCLI nodeAttributesCLI;
|
||||
private ByteArrayOutputStream errOutBytes = new ByteArrayOutputStream();
|
||||
private String errOutput;
|
||||
|
||||
@Before
|
||||
public void configure() throws IOException, YarnException {
|
||||
admin = mock(ResourceManagerAdministrationProtocol.class);
|
||||
|
||||
when(admin.mapAttributesToNodes(any(NodesToAttributesMappingRequest.class)))
|
||||
.thenAnswer(new Answer<NodesToAttributesMappingResponse>() {
|
||||
@Override
|
||||
public NodesToAttributesMappingResponse answer(
|
||||
InvocationOnMock invocation) throws Throwable {
|
||||
request =
|
||||
(NodesToAttributesMappingRequest) invocation.getArguments()[0];
|
||||
return NodesToAttributesMappingResponse.newInstance();
|
||||
}
|
||||
});
|
||||
|
||||
nodeAttributesCLI = new NodeAttributesCLI(new Configuration()) {
|
||||
@Override
|
||||
protected ResourceManagerAdministrationProtocol createAdminProtocol()
|
||||
throws IOException {
|
||||
return admin;
|
||||
}
|
||||
};
|
||||
|
||||
nodeAttributesCLI.setErrOut(new PrintStream(errOutBytes));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testHelp() throws Exception {
|
||||
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"
|
||||
+ " 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" };
|
||||
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"
|
||||
+ " information at the ResourceManager");
|
||||
|
||||
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"
|
||||
+ " 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" };
|
||||
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.");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testReplace() throws Exception {
|
||||
// --------------------------------
|
||||
// failure scenarios
|
||||
// --------------------------------
|
||||
// parenthesis not match
|
||||
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" };
|
||||
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" };
|
||||
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)" };
|
||||
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)" };
|
||||
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" };
|
||||
assertTrue("Should fail as no attribute mappings specified",
|
||||
0 != runTool(args));
|
||||
assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG);
|
||||
|
||||
// no labels, should fail
|
||||
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);
|
||||
|
||||
// no labels, should fail
|
||||
args = new String[] { "-replace", " " };
|
||||
assertTrue(0 != runTool(args));
|
||||
assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG);
|
||||
|
||||
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" };
|
||||
assertTrue("Should not fail as attribute has been properly mapped",
|
||||
0 == runTool(args));
|
||||
List<NodeToAttributes> nodeAttributesList = new ArrayList<>();
|
||||
List<NodeAttribute> attributes = new ArrayList<>();
|
||||
attributes.add(
|
||||
NodeAttribute.newInstance("key", NodeAttributeType.STRING, "value"));
|
||||
attributes.add(
|
||||
NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "val2"));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes));
|
||||
|
||||
// for node y
|
||||
attributes = new ArrayList<>();
|
||||
attributes.add(
|
||||
NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "val23"));
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key3", NodeAttributeType.STRING, ""));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("y", attributes));
|
||||
|
||||
// for node y
|
||||
attributes = new ArrayList<>();
|
||||
attributes.add(
|
||||
NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "val23"));
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key3", NodeAttributeType.STRING, ""));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("y", attributes));
|
||||
|
||||
// for node z
|
||||
attributes = new ArrayList<>();
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, ""));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes));
|
||||
|
||||
NodesToAttributesMappingRequest expected =
|
||||
NodesToAttributesMappingRequest.newInstance(
|
||||
AttributeMappingOperationType.REPLACE, nodeAttributesList, false);
|
||||
assertTrue(request.equals(expected));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testRemove() throws Exception {
|
||||
// --------------------------------
|
||||
// failure scenarios
|
||||
// --------------------------------
|
||||
// parenthesis not match
|
||||
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");
|
||||
// --------------------------------
|
||||
// success scenarios
|
||||
// --------------------------------
|
||||
args =
|
||||
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<>();
|
||||
List<NodeAttribute> attributes = new ArrayList<>();
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key2", NodeAttributeType.STRING, ""));
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key3", NodeAttributeType.STRING, ""));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes));
|
||||
|
||||
// for node z
|
||||
attributes = new ArrayList<>();
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, ""));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes));
|
||||
|
||||
NodesToAttributesMappingRequest expected =
|
||||
NodesToAttributesMappingRequest.newInstance(
|
||||
AttributeMappingOperationType.REMOVE, nodeAttributesList, true);
|
||||
assertTrue(request.equals(expected));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testAdd() throws Exception {
|
||||
// --------------------------------
|
||||
// failure scenarios
|
||||
// --------------------------------
|
||||
// parenthesis not match
|
||||
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");
|
||||
// --------------------------------
|
||||
// success scenarios
|
||||
// --------------------------------
|
||||
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<>();
|
||||
List<NodeAttribute> attributes = new ArrayList<>();
|
||||
attributes.add(
|
||||
NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "123"));
|
||||
attributes.add(
|
||||
NodeAttribute.newInstance("key3", NodeAttributeType.STRING, "abc"));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes));
|
||||
|
||||
// for node z
|
||||
attributes = new ArrayList<>();
|
||||
attributes
|
||||
.add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, ""));
|
||||
nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes));
|
||||
|
||||
NodesToAttributesMappingRequest expected =
|
||||
NodesToAttributesMappingRequest.newInstance(
|
||||
AttributeMappingOperationType.ADD, nodeAttributesList, true);
|
||||
assertTrue(request.equals(expected));
|
||||
}
|
||||
|
||||
private void assertFailureMessageContains(String... messages) {
|
||||
assertOutputContains(messages);
|
||||
assertOutputContains(NodeAttributesCLI.USAGE_YARN_NODE_ATTRIBUTES);
|
||||
}
|
||||
|
||||
private void assertOutputContains(String... messages) {
|
||||
for (String message : messages) {
|
||||
if (!errOutput.contains(message)) {
|
||||
fail("Expected output to contain '" + message
|
||||
+ "' but err_output was:\n" + errOutput);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private int runTool(String... args) throws Exception {
|
||||
errOutBytes.reset();
|
||||
LOG.info("Running: NodeAttributesCLI " + Joiner.on(" ").join(args));
|
||||
int ret = nodeAttributesCLI.run(args);
|
||||
errOutput = new String(errOutBytes.toByteArray(), Charsets.UTF_8);
|
||||
LOG.info("Err_output:\n" + errOutput);
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user