public static int getApplicableNodeCountForAM(RMContext rmContext,
      Configuration conf, List<ResourceRequest> amReqs) {
    // Determine the list of nodes that are eligible based on the strict
    // resource requests
    Set<NodeId> nodesForReqs = new HashSet<>();
    for (ResourceRequest amReq : amReqs) {
      if (amReq.getRelaxLocality() &&
          !amReq.getResourceName().equals(ResourceRequest.ANY)) {
        nodesForReqs.addAll(
            rmContext.getScheduler().getNodeIds(amReq.getResourceName()));
      }
    }

    if (YarnConfiguration.areNodeLabelsEnabled(conf)) {
      // Determine the list of nodes that are eligible based on the node label
      String amNodeLabelExpression = amReqs.get(0).getNodeLabelExpression();
      Set<NodeId> nodesForLabels =
          getNodeIdsForLabel(rmContext, amNodeLabelExpression);
      if (nodesForLabels != null && !nodesForLabels.isEmpty()) {
        // If only node labels, strip out any wildcard NodeIds and return
        if (nodesForReqs.isEmpty()) {
          for (Iterator<NodeId> it = nodesForLabels.iterator(); it.hasNext();) {
            if (it.next().getPort() == 0) {
              it.remove();
            }
          }
          return nodesForLabels.size();
        } else {
          // The NodeIds common to both the strict resource requests and the
          // node label is the eligible set
          return Sets.intersection(nodesForReqs, nodesForLabels).size();
        }
      }
    }

    // If no strict resource request NodeIds nor node label NodeIds, then just
    // return the entire cluster
    if (nodesForReqs.isEmpty()) {
      return rmContext.getScheduler().getNumClusterNodes();
    }
    // No node label NodeIds, so return the strict resource request NodeIds
    return nodesForReqs.size();
  }
@Test
  public void testGetApplicableNodeCountForAMLabels() throws Exception {
    Set<NodeId> noLabelNodes = new HashSet<>();
    for (int i = 0; i < 80; i++) {
      noLabelNodes.add(NodeId.newInstance("host" + i, 1234));
    }
    Set<NodeId> label1Nodes = new HashSet<>();
    for (int i = 80; i < 90; i++) {
      label1Nodes.add(NodeId.newInstance("host" + i, 1234));
    }
    label1Nodes.add(NodeId.newInstance("host101", 0));
    label1Nodes.add(NodeId.newInstance("host102", 0));
    Map<String, Set<NodeId>> label1NodesMap = new HashMap<>();
    label1NodesMap.put("label1", label1Nodes);

    YarnConfiguration conf = new YarnConfiguration();
    conf.setBoolean(YarnConfiguration.NODE_LABELS_ENABLED, true);
    ResourceScheduler scheduler = Mockito.mock(ResourceScheduler.class);
    Mockito.when(scheduler.getNumClusterNodes()).thenReturn(100);
    RMContext rmContext = Mockito.mock(RMContext.class);
    Mockito.when(rmContext.getScheduler()).thenReturn(scheduler);
    RMNodeLabelsManager labMan = Mockito.mock(RMNodeLabelsManager.class);
    Mockito.when(labMan.getNodesWithoutALabel()).thenReturn(noLabelNodes);
    Mockito.when(labMan.getLabelsToNodes(Collections.singleton("label1")))
        .thenReturn(label1NodesMap);
    Mockito.when(rmContext.getNodeLabelManager()).thenReturn(labMan);

    ResourceRequest anyReq = createResourceRequest(ResourceRequest.ANY,
        true, null);
    List<ResourceRequest> reqs = new ArrayList<>();
    reqs.add(anyReq);
    Assert.assertEquals(80,
        RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
    anyReq.setNodeLabelExpression("label1");
    Assert.assertEquals(10,
        RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  }