@SuppressWarnings("unchecked")
public void addToCluserNodeLabels(Collection<NodeLabel> labels) throws IOException {
if (!nodeLabelsEnabled) {
LOG.error(NODE_LABELS_NOT_ENABLED_ERR);
throw new IOException(NODE_LABELS_NOT_ENABLED_ERR);
}
if (null == labels || labels.isEmpty()) {
return;
}
List<NodeLabel> newLabels = new ArrayList<NodeLabel>();
normalizeNodeLabels(labels);
// do a check before actual adding them, will throw exception if any of them
// doesn't meet label name requirement
for (NodeLabel label : labels) {
checkAndThrowLabelName(label.getName());
}
for (NodeLabel label : labels) {
// shouldn't overwrite it to avoid changing the Label.resource
if (this.labelCollections.get(label.getName()) == null) {
this.labelCollections.put(label.getName(), new RMNodeLabel(label));
newLabels.add(label);
}
}
if (null != dispatcher && !newLabels.isEmpty()) {
dispatcher.getEventHandler().handle(new StoreNewClusterNodeLabels(newLabels));
}
LOG.info("Add labels: [" + StringUtils.join(labels.iterator(), ",") + "]");
}
/**
* Add multiple node labels to repository
*
* @param labels new node labels added
*/
@VisibleForTesting
public void addToCluserNodeLabelsWithDefaultExclusivity(Set<String> labels) throws IOException {
Set<NodeLabel> nodeLabels = new HashSet<NodeLabel>();
for (String label : labels) {
nodeLabels.add(NodeLabel.newInstance(label));
}
addToCluserNodeLabels(nodeLabels);
}
public List<NodeLabel> getClusterNodeLabels() {
try {
readLock.lock();
List<NodeLabel> nodeLabels = new ArrayList<>();
for (RMNodeLabel label : labelCollections.values()) {
if (!label.getLabelName().equals(NO_LABEL)) {
nodeLabels.add(NodeLabel.newInstance(label.getLabelName(), label.getIsExclusive()));
}
}
return nodeLabels;
} finally {
readLock.unlock();
}
}
@Test
public void testNodeRegistrationWithLabels() throws Exception {
writeToHostsFile("host2");
Configuration conf = new Configuration();
conf.set(YarnConfiguration.RM_NODES_INCLUDE_FILE_PATH, hostFile.getAbsolutePath());
conf.set(
YarnConfiguration.NODELABEL_CONFIGURATION_TYPE,
YarnConfiguration.DISTRIBUTED_NODELABEL_CONFIGURATION_TYPE);
final RMNodeLabelsManager nodeLabelsMgr = new NullRMNodeLabelsManager();
rm =
new MockRM(conf) {
@Override
protected RMNodeLabelsManager createNodeLabelManager() {
return nodeLabelsMgr;
}
};
rm.start();
try {
nodeLabelsMgr.addToCluserNodeLabelsWithDefaultExclusivity(toSet("A", "B", "C"));
} catch (IOException e) {
Assert.fail("Caught Exception while intializing");
e.printStackTrace();
}
ResourceTrackerService resourceTrackerService = rm.getResourceTrackerService();
RegisterNodeManagerRequest registerReq = Records.newRecord(RegisterNodeManagerRequest.class);
NodeId nodeId = NodeId.newInstance("host2", 1234);
Resource capability = BuilderUtils.newResource(1024, 1);
registerReq.setResource(capability);
registerReq.setNodeId(nodeId);
registerReq.setHttpPort(1234);
registerReq.setNMVersion(YarnVersionInfo.getVersion());
registerReq.setNodeLabels(toSet(NodeLabel.newInstance("A")));
RegisterNodeManagerResponse response = resourceTrackerService.registerNodeManager(registerReq);
Assert.assertEquals(
"Action should be normal on valid Node Labels",
NodeAction.NORMAL,
response.getNodeAction());
assertCollectionEquals(
nodeLabelsMgr.getNodeLabels().get(nodeId),
ResourceTrackerService.convertToStringSet(registerReq.getNodeLabels()));
Assert.assertTrue(
"Valid Node Labels were not accepted by RM", response.getAreNodeLabelsAcceptedByRM());
rm.stop();
}
private void normalizeNodeLabels(Collection<NodeLabel> labels) {
for (NodeLabel label : labels) {
label.setName(normalizeLabel(label.getName()));
}
}