org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerNode Java Examples

@Test
public void testAddAndRemoveAppFromCapacityScheduler() throws Exception {
  CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
  setupQueueConfiguration(conf);
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> cs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();

  SchedulerApplication<SchedulerApplicationAttempt> app =
      TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
        cs.getSchedulerApplications(), cs, "a1");
  Assert.assertEquals("a1", app.getQueue().getQueueName());
}
 

@Test
public void testAddAndRemoveAppFromCapacityScheduler() throws Exception {
  CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
  setupQueueConfiguration(conf);
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> cs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();

  SchedulerApplication<SchedulerApplicationAttempt> app =
      TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
        cs.getSchedulerApplications(), cs, "a1");
  Assert.assertEquals("a1", app.getQueue().getQueueName());
}
 

@SuppressWarnings("unchecked")
@Override
public void serviceStart() throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).start();
  super.serviceStart();
}
 

@Before
public void setUp() throws Exception {
  NodeStore store = new NodeStore();
  NodeIdProto nodeId = NodeIdProto.newBuilder().setHost("localhost").setPort(8000).build();
  RMNode rmNode = new RMNodeImpl(new NodeIdPBImpl(nodeId), new MockRMContext(), "localhost", 8000, 8070, new NodeBase(),
          new ResourcePBImpl(), "1.0");
  SchedulerNode node = new FiCaSchedulerNode(rmNode, false);
  store.add(node);
  manager = new OfferLifecycleManager(store, new MyriadDriver(new MockSchedulerDriver()));
}
 

@Override
public void afterSchedulerEventHandled(SchedulerEvent event) {
  switch (event.getType()) {
    case NODE_ADDED:
      if (!(event instanceof NodeAddedSchedulerEvent)) {
        LOGGER.error("{} not an instance of {}", event.getClass().getName(), NodeAddedSchedulerEvent.class.getName());
        return;
      }

      NodeAddedSchedulerEvent nodeAddedEvent = (NodeAddedSchedulerEvent) event;
      NodeId nodeId = nodeAddedEvent.getAddedRMNode().getNodeID();
      String host = nodeId.getHost();

      SchedulerNode node = yarnScheduler.getSchedulerNode(nodeId);
      nodeStore.add(node);
      LOGGER.info("afterSchedulerEventHandled: NM registration from node {}", host);
      break;

    case NODE_UPDATE:
      if (!(event instanceof NodeUpdateSchedulerEvent)) {
        LOGGER.error("{} not an instance of {}", event.getClass().getName(), NodeUpdateSchedulerEvent.class.getName());
        return;
      }

      RMNode rmNode = ((NodeUpdateSchedulerEvent) event).getRMNode();
      handleContainerAllocation(rmNode);

      break;

    default:
      break;
  }
}
 

@SuppressWarnings("unchecked")
@Override
public void serviceStop() throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).stop();
  super.serviceStop();
}
 

@SuppressWarnings("unchecked")
@Override
public void serviceInit(Configuration conf) throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).init(conf);
  super.serviceInit(conf);
}
 

@Test
public void testAddAndRemoveAppFromFairScheduler() throws Exception {
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> scheduler =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) resourceManager
        .getResourceScheduler();
  TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
    scheduler.getSchedulerApplications(), scheduler, "default");
}
 

@Test
public void testAddAndRemoveAppFromFiFoScheduler() throws Exception {
  Configuration conf = new Configuration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> fs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();
  TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
    fs.getSchedulerApplications(), fs, "queue");
}
 

@SuppressWarnings("unchecked")
@Override
public void serviceStop() throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).stop();
  super.serviceStop();
}
 

@SuppressWarnings("unchecked")
@Override
public void serviceStart() throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).start();
  super.serviceStart();
}
 

@SuppressWarnings("unchecked")
@Override
public void serviceInit(Configuration conf) throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).init(conf);
  super.serviceInit(conf);
}
 

@Test
public void testAddAndRemoveAppFromFairScheduler() throws Exception {
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> scheduler =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) resourceManager
        .getResourceScheduler();
  TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
    scheduler.getSchedulerApplications(), scheduler, "default");
}
 

@Test
public void testAddAndRemoveAppFromFiFoScheduler() throws Exception {
  Configuration conf = new Configuration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> fs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();
  TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
    fs.getSchedulerApplications(), fs, "queue");
}
 

@SuppressWarnings("static-access")
public void clearQueueMetrics(RMApp app) {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) getResourceScheduler())
    .getSchedulerApplications().get(app.getApplicationId()).getQueue()
    .getMetrics().clearQueueMetrics();
}
 

@SuppressWarnings("static-access")
public void clearQueueMetrics(RMApp app) {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) getResourceScheduler())
    .getSchedulerApplications().get(app.getApplicationId()).getQueue()
    .getMetrics().clearQueueMetrics();
}
 

/**
 * Sort the given list of tasks by the number of App Master containers running on the corresponding NM node.
 *
 * @param taskIDs
 */
@Override
public void apply(List<Protos.TaskID> taskIDs) {
  if (LOGGER.isDebugEnabled()) {
    for (SchedulerNode node : schedulerNodes.values()) {
      LOGGER.debug("Host {} is running {} containers including {} App Masters", node.getNodeID().getHost(),
          node.getRunningContainers().size(), getNumAMContainers(node.getRunningContainers()));
    }
  }
  // We need to lock the YARN scheduler here. If we don't do that, then the YARN scheduler can
  // process HBs from NodeManagers and the state of SchedulerNode objects might change while we
  // are in the middle of sorting them based on the least number of AM containers.
  synchronized (yarnScheduler) {
    Collections.sort(taskIDs, new Comparator<Protos.TaskID>() {
      @Override
      public int compare(Protos.TaskID t1, Protos.TaskID t2) {
        SchedulerNode o1 = schedulerNodes.get(schedulerState.getTask(t1).getHostname());
        SchedulerNode o2 = schedulerNodes.get(schedulerState.getTask(t2).getHostname());

        if (o1 == null) { // a NM was launched by Myriad, but it hasn't yet registered with RM
          if (o2 == null) {
            return 0;
          } else {
            return -1;
          }
        } else if (o2 == null) {
          return 1;
        } // else, both the NMs have registered with RM

        List<RMContainer> runningContainers1 = o1.getRunningContainers();
        List<RMContainer> runningContainers2 = o2.getRunningContainers();

        Integer numRunningAMs1 = getNumAMContainers(runningContainers1);
        Integer numRunningAMs2 = getNumAMContainers(runningContainers2);

        Integer numRunningContainers1 = runningContainers1.size();
        Integer numRunningContainers2 = runningContainers2.size();

        // If two NMs are running equal number of AMs, sort them based on total num of running containers
        if (numRunningAMs1.compareTo(numRunningAMs2) == 0) {
          return numRunningContainers1.compareTo(numRunningContainers2);
        }
        return numRunningAMs1.compareTo(numRunningAMs2);
      }
    });
  }
}
 

private void onNodeRemoved(NodeRemovedSchedulerEvent event) {
  SchedulerNode schedulerNode = schedulerNodes.get(event.getRemovedRMNode().getNodeID().getHost());
  if (schedulerNode != null && schedulerNode.getNodeID().equals(event.getRemovedRMNode().getNodeID())) {
    schedulerNodes.remove(schedulerNode.getNodeID().getHost());
  }
}
 

private String getKey(SchedulerNode schedNode) {
  return schedNode.getNodeID().getHost();
}
 

public void add(SchedulerNode schedNode) {
  nodeMap.put(getKey(schedNode), new Node(schedNode));
}
 

public Node(SchedulerNode node) {
  this.node = node;
}
 

public SchedulerNode getNode() {
  return node;
}
 

public static SchedulerNode getSchedulerNode(String host, int port, int cores, int memory) {
  RMNode node = TestObjectFactory.getRMNode(host, port, cores, memory);
  return new FiCaSchedulerNode(node, false);
}
 

/**
 * Called whenever a NM HBs to RM. The NM's updates will already be recorded in the
 * SchedulerNode before this method is called.
 *
 * @param event
 */
private void onNodeUpdated(NodeUpdateSchedulerEvent event) {
  NodeId nodeID = event.getRMNode().getNodeID();
  SchedulerNode schedulerNode = yarnScheduler.getSchedulerNode(nodeID);
  schedulerNodes.put(nodeID.getHost(), schedulerNode); // keep track of only one node per host
}