org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager Java Examples

private static void normalizeNodeLabelExpressionInRequest(
    ResourceRequest resReq, QueueInfo queueInfo) {

  String labelExp = resReq.getNodeLabelExpression();

  // if queue has default label expression, and RR doesn't have, use the
  // default label expression of queue
  if (labelExp == null && queueInfo != null && ResourceRequest.ANY
      .equals(resReq.getResourceName())) {
    labelExp = queueInfo.getDefaultNodeLabelExpression();
  }

  // If labelExp still equals to null, set it to be NO_LABEL
  if (labelExp == null) {
    labelExp = RMNodeLabelsManager.NO_LABEL;
  }
  resReq.setNodeLabelExpression(labelExp);
}
 

@Before
public void setUp() throws Exception {
  resourceManager = new ResourceManager() {
    @Override
    protected RMNodeLabelsManager createNodeLabelManager() {
      RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
      mgr.init(getConfig());
      return mgr;
    }
  };
  CapacitySchedulerConfiguration csConf 
     = new CapacitySchedulerConfiguration();
  setupQueueConfiguration(csConf);
  YarnConfiguration conf = new YarnConfiguration(csConf);
  conf.setClass(YarnConfiguration.RM_SCHEDULER, 
      CapacityScheduler.class, ResourceScheduler.class);
  resourceManager.init(conf);
  resourceManager.getRMContext().getContainerTokenSecretManager().rollMasterKey();
  resourceManager.getRMContext().getNMTokenSecretManager().rollMasterKey();
  ((AsyncDispatcher)resourceManager.getRMContext().getDispatcher()).start();
  mockContext = mock(RMContext.class);
  when(mockContext.getConfigurationProvider()).thenReturn(
      new LocalConfigurationProvider());
}
 

@Test
public void testQueueParsingWhenLabelsNotExist() throws IOException {
  YarnConfiguration conf = new YarnConfiguration();
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration(conf);
  setupQueueConfigurationWithLabels(csConf);

  CapacityScheduler capacityScheduler = new CapacityScheduler();
  RMContextImpl rmContext =
      new RMContextImpl(null, null, null, null, null, null,
          new RMContainerTokenSecretManager(csConf),
          new NMTokenSecretManagerInRM(csConf),
          new ClientToAMTokenSecretManagerInRM(), null);
  
  RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager();
  nodeLabelsManager.init(conf);
  nodeLabelsManager.start();
  
  rmContext.setNodeLabelManager(nodeLabelsManager);
  capacityScheduler.setConf(csConf);
  capacityScheduler.setRMContext(rmContext);
  capacityScheduler.init(csConf);
  capacityScheduler.start();
  ServiceOperations.stopQuietly(capacityScheduler);
  ServiceOperations.stopQuietly(nodeLabelsManager);
}
 

public static void checkIfLabelInClusterNodeLabels(RMNodeLabelsManager mgr,
    Set<String> labels) throws IOException {
  if (mgr == null) {
    if (labels != null && !labels.isEmpty()) {
      throw new IOException("NodeLabelManager is null, please check");
    }
    return;
  }

  if (labels != null) {
    for (String label : labels) {
      if (!label.equals(RMNodeLabelsManager.ANY)
          && !mgr.containsNodeLabel(label)) {
        throw new IOException("NodeLabelManager doesn't include label = "
            + label + ", please check.");
      }
    }
  }
}
 

public static boolean checkQueueAccessToNode(Set<String> queueLabels,
    Set<String> nodeLabels) {
  // if queue's label is *, it can access any node
  if (queueLabels != null && queueLabels.contains(RMNodeLabelsManager.ANY)) {
    return true;
  }
  // any queue can access to a node without label
  if (nodeLabels == null || nodeLabels.isEmpty()) {
    return true;
  }
  // a queue can access to a node only if it contains any label of the node
  if (queueLabels != null
      && Sets.intersection(queueLabels, nodeLabels).size() > 0) {
    return true;
  }
  // sorry, you cannot access
  return false;
}
 

private void updateCurrentResourceLimits(
    ResourceLimits currentResourceLimits, Resource clusterResource) {
  // TODO: need consider non-empty node labels when resource limits supports
  // node labels
  // Even if ParentQueue will set limits respect child's max queue capacity,
  // but when allocating reserved container, CapacityScheduler doesn't do
  // this. So need cap limits by queue's max capacity here.
  this.cachedResourceLimitsForHeadroom = new ResourceLimits(currentResourceLimits.getLimit());
  Resource queueMaxResource =
      Resources.multiplyAndNormalizeDown(resourceCalculator, labelManager
          .getResourceByLabel(RMNodeLabelsManager.NO_LABEL, clusterResource),
          queueCapacities
              .getAbsoluteMaximumCapacity(RMNodeLabelsManager.NO_LABEL),
          minimumAllocation);
  this.cachedResourceLimitsForHeadroom.setLimit(Resources.min(resourceCalculator,
      clusterResource, queueMaxResource, currentResourceLimits.getLimit()));
}
 

synchronized void allocateResource(Resource clusterResource, 
    Resource resource, Set<String> nodeLabels) {
  
  // Update usedResources by labels
  if (nodeLabels == null || nodeLabels.isEmpty()) {
    queueUsage.incUsed(resource);
  } else {
    Set<String> anls = (accessibleLabels.contains(RMNodeLabelsManager.ANY))
        ? labelManager.getClusterNodeLabels() : accessibleLabels;
    for (String label : Sets.intersection(anls, nodeLabels)) {
      queueUsage.incUsed(label, resource);
    }
  }

  ++numContainers;
  if (null == nodeLabels || nodeLabels.isEmpty()) {
    CSQueueUtils.updateQueueStatistics(resourceCalculator, this, getParent(),
        labelManager.getResourceByLabel(RMNodeLabelsManager.NO_LABEL, clusterResource), minimumAllocation);
  }
}
 

private Resource getCurrentLimitResource(String nodeLabel,
    Resource clusterResource, ResourceLimits currentResourceLimits) {
  /*
   * Current limit resource: For labeled resource: limit = queue-max-resource
   * (TODO, this part need update when we support labeled-limit) For
   * non-labeled resource: limit = min(queue-max-resource,
   * limit-set-by-parent)
   */
  Resource queueMaxResource =
      Resources.multiplyAndNormalizeDown(resourceCalculator,
          labelManager.getResourceByLabel(nodeLabel, clusterResource),
          queueCapacities.getAbsoluteMaximumCapacity(nodeLabel), minimumAllocation);
  if (nodeLabel.equals(RMNodeLabelsManager.NO_LABEL)) {
    return Resources.min(resourceCalculator, clusterResource,
        queueMaxResource, currentResourceLimits.getLimit());
  }
  return queueMaxResource;
}
 

private void updateCurrentResourceLimits(
    ResourceLimits currentResourceLimits, Resource clusterResource) {
  // TODO: need consider non-empty node labels when resource limits supports
  // node labels
  // Even if ParentQueue will set limits respect child's max queue capacity,
  // but when allocating reserved container, CapacityScheduler doesn't do
  // this. So need cap limits by queue's max capacity here.
  this.cachedResourceLimitsForHeadroom = new ResourceLimits(currentResourceLimits.getLimit());
  Resource queueMaxResource =
      Resources.multiplyAndNormalizeDown(resourceCalculator, labelManager
          .getResourceByLabel(RMNodeLabelsManager.NO_LABEL, clusterResource),
          queueCapacities
              .getAbsoluteMaximumCapacity(RMNodeLabelsManager.NO_LABEL),
          minimumAllocation);
  this.cachedResourceLimitsForHeadroom.setLimit(Resources.min(resourceCalculator,
      clusterResource, queueMaxResource, currentResourceLimits.getLimit()));
}
 

public static void checkIfLabelInClusterNodeLabels(RMNodeLabelsManager mgr,
    Set<String> labels) throws IOException {
  if (mgr == null) {
    if (labels != null && !labels.isEmpty()) {
      throw new IOException("NodeLabelManager is null, please check");
    }
    return;
  }

  if (labels != null) {
    for (String label : labels) {
      if (!label.equals(RMNodeLabelsManager.ANY)
          && !mgr.containsNodeLabel(label)) {
        throw new IOException("NodeLabelManager doesn't include label = "
            + label + ", please check.");
      }
    }
  }
}
 

@Test
public void testQueueParsingWhenLabelsNotExist() throws IOException {
  YarnConfiguration conf = new YarnConfiguration();
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration(conf);
  setupQueueConfigurationWithLabels(csConf);

  CapacityScheduler capacityScheduler = new CapacityScheduler();
  RMContextImpl rmContext =
      new RMContextImpl(null, null, null, null, null, null,
          new RMContainerTokenSecretManager(csConf),
          new NMTokenSecretManagerInRM(csConf),
          new ClientToAMTokenSecretManagerInRM(), null);
  
  RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager();
  nodeLabelsManager.init(conf);
  nodeLabelsManager.start();
  
  rmContext.setNodeLabelManager(nodeLabelsManager);
  capacityScheduler.setConf(csConf);
  capacityScheduler.setRMContext(rmContext);
  capacityScheduler.init(csConf);
  capacityScheduler.start();
  ServiceOperations.stopQuietly(capacityScheduler);
  ServiceOperations.stopQuietly(nodeLabelsManager);
}
 

@GET
@Path("/scheduler")
@Produces({ MediaType.APPLICATION_JSON, MediaType.APPLICATION_XML })
public SchedulerTypeInfo getSchedulerInfo() {
  init();
  ResourceScheduler rs = rm.getResourceScheduler();
  SchedulerInfo sinfo;
  if (rs instanceof CapacityScheduler) {
    CapacityScheduler cs = (CapacityScheduler) rs;
    CSQueue root = cs.getRootQueue();
    sinfo =
        new CapacitySchedulerInfo(root, cs, new NodeLabel(
            RMNodeLabelsManager.NO_LABEL));
  } else if (rs instanceof FairScheduler) {
    FairScheduler fs = (FairScheduler) rs;
    sinfo = new FairSchedulerInfo(fs);
  } else if (rs instanceof FifoScheduler) {
    sinfo = new FifoSchedulerInfo(this.rm);
  } else {
    throw new NotFoundException("Unknown scheduler configured");
  }
  return new SchedulerTypeInfo(sinfo);
}
 

@Override
public GetClusterNodeLabelsResponse getClusterNodeLabels(
    GetClusterNodeLabelsRequest request) throws YarnException, IOException {
  RMNodeLabelsManager labelsMgr = rmContext.getNodeLabelManager();
  GetClusterNodeLabelsResponse response =
      GetClusterNodeLabelsResponse.newInstance(
          labelsMgr.getClusterNodeLabels());
  return response;
}
 

/**
 * Get configured node labels in a given queuePath
 */
public Set<String> getConfiguredNodeLabels(String queuePath) {
  Set<String> configuredNodeLabels = new HashSet<String>();
  Entry<String, String> e = null;
  
  Iterator<Entry<String, String>> iter = iterator();
  while (iter.hasNext()) {
    e = iter.next();
    String key = e.getKey();

    if (key.startsWith(getQueuePrefix(queuePath) + ACCESSIBLE_NODE_LABELS
        + DOT)) {
      // Find <label-name> in
      // <queue-path>.accessible-node-labels.<label-name>.property
      int labelStartIdx =
          key.indexOf(ACCESSIBLE_NODE_LABELS)
              + ACCESSIBLE_NODE_LABELS.length() + 1;
      int labelEndIndx = key.indexOf('.', labelStartIdx);
      String labelName = key.substring(labelStartIdx, labelEndIndx);
      configuredNodeLabels.add(labelName);
    }
  }
  
  // always add NO_LABEL
  configuredNodeLabels.add(RMNodeLabelsManager.NO_LABEL);
  
  return configuredNodeLabels;
}
 

public Set<String> getAccessibleNodeLabels(String queue) {
  String accessibleLabelStr =
      get(getQueuePrefix(queue) + ACCESSIBLE_NODE_LABELS);

  // When accessible-label is null, 
  if (accessibleLabelStr == null) {
    // Only return null when queue is not ROOT
    if (!queue.equals(ROOT)) {
      return null;
    }
  } else {
    // print a warning when accessibleNodeLabel specified in config and queue
    // is ROOT
    if (queue.equals(ROOT)) {
      LOG.warn("Accessible node labels for root queue will be ignored,"
          + " it will be automatically set to \"*\".");
    }
  }

  // always return ANY for queue root
  if (queue.equals(ROOT)) {
    return ImmutableSet.of(RMNodeLabelsManager.ANY);
  }

  // In other cases, split the accessibleLabelStr by ","
  Set<String> set = new HashSet<String>();
  for (String str : accessibleLabelStr.split(",")) {
    if (!str.trim().isEmpty()) {
      set.add(str.trim());
    }
  }
  
  // if labels contains "*", only keep ANY behind
  if (set.contains(RMNodeLabelsManager.ANY)) {
    set.clear();
    set.add(RMNodeLabelsManager.ANY);
  }
  return Collections.unmodifiableSet(set);
}
 

private void initializeNodeLabels() throws IOException {
  RMContext rmContext = distShellTest.yarnCluster.getResourceManager(0).getRMContext();

  // Setup node labels
  RMNodeLabelsManager labelsMgr = rmContext.getNodeLabelManager();
  Set<String> labels = new HashSet<String>();
  labels.add("x");
  labelsMgr.addToCluserNodeLabels(labels);

  // Setup queue access to node labels
  distShellTest.conf.set("yarn.scheduler.capacity.root.accessible-node-labels", "x");
  distShellTest.conf.set("yarn.scheduler.capacity.root.accessible-node-labels.x.capacity",
      "100");
  distShellTest.conf.set("yarn.scheduler.capacity.root.default.accessible-node-labels", "x");
  distShellTest.conf.set(
      "yarn.scheduler.capacity.root.default.accessible-node-labels.x.capacity",
      "100");

  rmContext.getScheduler().reinitialize(distShellTest.conf, rmContext);

  // Fetch node-ids from yarn cluster
  NodeId[] nodeIds = new NodeId[NUM_NMS];
  for (int i = 0; i < NUM_NMS; i++) {
    NodeManager mgr = distShellTest.yarnCluster.getNodeManager(i);
    nodeIds[i] = mgr.getNMContext().getNodeId();
  }

  // Set label x to NM[1]
  labelsMgr.addLabelsToNode(ImmutableMap.of(nodeIds[1], labels));
}
 

public Set<String> getAccessibleNodeLabels(String queue) {
  String accessibleLabelStr =
      get(getQueuePrefix(queue) + ACCESSIBLE_NODE_LABELS);

  // When accessible-label is null, 
  if (accessibleLabelStr == null) {
    // Only return null when queue is not ROOT
    if (!queue.equals(ROOT)) {
      return null;
    }
  } else {
    // print a warning when accessibleNodeLabel specified in config and queue
    // is ROOT
    if (queue.equals(ROOT)) {
      LOG.warn("Accessible node labels for root queue will be ignored,"
          + " it will be automatically set to \"*\".");
    }
  }

  // always return ANY for queue root
  if (queue.equals(ROOT)) {
    return ImmutableSet.of(RMNodeLabelsManager.ANY);
  }

  // In other cases, split the accessibleLabelStr by ","
  Set<String> set = new HashSet<String>();
  for (String str : accessibleLabelStr.split(",")) {
    if (!str.trim().isEmpty()) {
      set.add(str.trim());
    }
  }
  
  // if labels contains "*", only keep ANY behind
  if (set.contains(RMNodeLabelsManager.ANY)) {
    set.clear();
    set.add(RMNodeLabelsManager.ANY);
  }
  return Collections.unmodifiableSet(set);
}
 

@Override
public GetClusterNodeLabelsResponse getClusterNodeLabels(
    GetClusterNodeLabelsRequest request) throws YarnException, IOException {
  RMNodeLabelsManager labelsMgr = rmContext.getNodeLabelManager();
  GetClusterNodeLabelsResponse response =
      GetClusterNodeLabelsResponse.newInstance(
          labelsMgr.getClusterNodeLabels());
  return response;
}
 

/**
 * Gets the labels which are accessible by this queue. If ANY label can be
 * accessed, put all labels in the set.
 * @return accessiglbe node labels
 */
protected final Set<String> getAccessibleLabelSet() {
  Set<String> nodeLabels = new HashSet<String>();
  if (this.getAccessibleNodeLabels().contains(RMNodeLabelsManager.ANY)) {
    nodeLabels.addAll(labelManager.getClusterNodeLabels());
  } else {
    nodeLabels.addAll(this.getAccessibleNodeLabels());
  }
  nodeLabels.add(RMNodeLabelsManager.NO_LABEL);
  return nodeLabels;
}
 

@Test
public void testQueueParsingWhenLabelsNotExistedInNodeLabelManager()
    throws IOException {
  YarnConfiguration conf = new YarnConfiguration();
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration(conf);
  setupQueueConfigurationWithLabels(csConf);

  CapacityScheduler capacityScheduler = new CapacityScheduler();
  RMContextImpl rmContext =
      new RMContextImpl(null, null, null, null, null, null,
          new RMContainerTokenSecretManager(csConf),
          new NMTokenSecretManagerInRM(csConf),
          new ClientToAMTokenSecretManagerInRM(), null);
  
  RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager();
  nodeLabelsManager.init(conf);
  nodeLabelsManager.start();
  
  rmContext.setNodeLabelManager(nodeLabelsManager);
  capacityScheduler.setConf(csConf);
  capacityScheduler.setRMContext(rmContext);
  capacityScheduler.init(csConf);
  capacityScheduler.start();
  ServiceOperations.stopQuietly(capacityScheduler);
  ServiceOperations.stopQuietly(nodeLabelsManager);
}
 

@Override
protected RMNodeLabelsManager createNodeLabelManager()
    throws InstantiationException, IllegalAccessException {
  if (useNullRMNodeLabelsManager) {
    RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
    mgr.init(getConfig());
    return mgr;
  } else {
    return super.createNodeLabelManager();
  }
}
 

@Test
public void testQueueParsingWhenLabelsInheritedNotExistedInNodeLabelManager()
    throws IOException {
  YarnConfiguration conf = new YarnConfiguration();
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration(conf);
  setupQueueConfigurationWithLabelsInherit(csConf);

  CapacityScheduler capacityScheduler = new CapacityScheduler();
  RMContextImpl rmContext =
      new RMContextImpl(null, null, null, null, null, null,
          new RMContainerTokenSecretManager(csConf),
          new NMTokenSecretManagerInRM(csConf),
          new ClientToAMTokenSecretManagerInRM(), null);
  
  RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager();
  nodeLabelsManager.init(conf);
  nodeLabelsManager.start();
  
  rmContext.setNodeLabelManager(nodeLabelsManager);
  capacityScheduler.setConf(csConf);
  capacityScheduler.setRMContext(rmContext);
  capacityScheduler.init(csConf);
  capacityScheduler.start();
  ServiceOperations.stopQuietly(capacityScheduler);
  ServiceOperations.stopQuietly(nodeLabelsManager);
}
 

@Test
public void testSingleLevelQueueParsingWhenLabelsNotExistedInNodeLabelManager()
    throws IOException {
  YarnConfiguration conf = new YarnConfiguration();
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration(conf);
  setupQueueConfigurationWithSingleLevel(csConf);

  CapacityScheduler capacityScheduler = new CapacityScheduler();
  RMContextImpl rmContext =
      new RMContextImpl(null, null, null, null, null, null,
          new RMContainerTokenSecretManager(csConf),
          new NMTokenSecretManagerInRM(csConf),
          new ClientToAMTokenSecretManagerInRM(), null);
  
  RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager();
  nodeLabelsManager.init(conf);
  nodeLabelsManager.start();
  
  rmContext.setNodeLabelManager(nodeLabelsManager);
  capacityScheduler.setConf(csConf);
  capacityScheduler.setRMContext(rmContext);
  capacityScheduler.init(csConf);
  capacityScheduler.start();
  ServiceOperations.stopQuietly(capacityScheduler);
  ServiceOperations.stopQuietly(nodeLabelsManager);
}
 

@Test
public void testQueueParsingWhenLabelsInheritedNotExistedInNodeLabelManager()
    throws IOException {
  YarnConfiguration conf = new YarnConfiguration();
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration(conf);
  setupQueueConfigurationWithLabelsInherit(csConf);

  CapacityScheduler capacityScheduler = new CapacityScheduler();
  RMContextImpl rmContext =
      new RMContextImpl(null, null, null, null, null, null,
          new RMContainerTokenSecretManager(csConf),
          new NMTokenSecretManagerInRM(csConf),
          new ClientToAMTokenSecretManagerInRM(), null);
  
  RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager();
  nodeLabelsManager.init(conf);
  nodeLabelsManager.start();
  
  rmContext.setNodeLabelManager(nodeLabelsManager);
  capacityScheduler.setConf(csConf);
  capacityScheduler.setRMContext(rmContext);
  capacityScheduler.init(csConf);
  capacityScheduler.start();
  ServiceOperations.stopQuietly(capacityScheduler);
  ServiceOperations.stopQuietly(nodeLabelsManager);
}
 

/**
 * This method returns all non labeled resources.
 * 
 * @param clusterResources
 * @return Resources
 */
private Resource getNonLabeledResources(Resource clusterResources) {
  RMContext rmcontext = scheduler.getRMContext();
  RMNodeLabelsManager lm = rmcontext.getNodeLabelManager();
  Resource res = lm.getResourceByLabel(RMNodeLabelsManager.NO_LABEL,
      clusterResources);
  return res == null ? clusterResources : res;
}
 

@Override
protected RMNodeLabelsManager createNodeLabelManager()
    throws InstantiationException, IllegalAccessException {
  if (useNullRMNodeLabelsManager) {
    RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
    mgr.init(getConfig());
    return mgr;
  } else {
    return super.createNodeLabelManager();
  }
}
 

private static RMContext getMockRMContext() {
  RMContext rmContext = mock(RMContext.class);
  RMNodeLabelsManager nlm = new NullRMNodeLabelsManager();
  nlm.init(new Configuration(false));
  when(rmContext.getNodeLabelManager()).thenReturn(nlm);
  return rmContext;
}
 

protected ResourceRequest createResourceRequest(
    int memory, int vcores, int gcores, String host, int priority, int numContainers,
    boolean relaxLocality) {
  ResourceRequest request = recordFactory.newRecordInstance(ResourceRequest.class);
  request.setCapability(BuilderUtils.newResource(memory, vcores, gcores));
  request.setResourceName(host);
  request.setNumContainers(numContainers);
  Priority prio = recordFactory.newRecordInstance(Priority.class);
  prio.setPriority(priority);
  request.setPriority(prio);
  request.setRelaxLocality(relaxLocality);
  request.setNodeLabelExpression(RMNodeLabelsManager.NO_LABEL);
  return request;
}
 

private static RMContext getMockRMContext() {
  RMContext rmContext = mock(RMContext.class);
  RMNodeLabelsManager nlm = new NullRMNodeLabelsManager();
  nlm.init(new Configuration(false));
  when(rmContext.getNodeLabelManager()).thenReturn(nlm);
  return rmContext;
}
 

private void checkPendingResourceGreaterThanZero(MockRM rm, String queueName,
    String label) {
  CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
  CSQueue queue = cs.getQueue(queueName);
  Assert.assertTrue(queue.getQueueResourceUsage()
      .getPending(label == null ? RMNodeLabelsManager.NO_LABEL : label)
      .getMemory() > 0);
}