Examples with AMRMClientAsync org.apache.hadoop.yarn.client.api.async.AMRMClientAsync used on opensource projects

Example 1 with AMRMClientAsync

use of org.apache.hadoop.yarn.client.api.async.AMRMClientAsync in project hadoop by apache.

the class TestDSAppMaster method testDSAppMasterAllocateHandler.

@SuppressWarnings("unchecked")
@Test
public void testDSAppMasterAllocateHandler() throws Exception {
    TestAppMaster master = new TestAppMaster();
    int targetContainers = 2;
    AMRMClientAsync mockClient = Mockito.mock(AMRMClientAsync.class);
    master.setAmRMClient(mockClient);
    master.setNumTotalContainers(targetContainers);
    Mockito.doNothing().when(mockClient).addContainerRequest(Matchers.any(AMRMClient.ContainerRequest.class));
    ApplicationMaster.RMCallbackHandler handler = master.getRMCallbackHandler();
    List<Container> containers = new ArrayList<>(1);
    ContainerId id1 = BuilderUtils.newContainerId(1, 1, 1, 1);
    containers.add(generateContainer(id1));
    master.numRequestedContainers.set(targetContainers);
    // first allocate a single container, everything should be fine
    handler.onContainersAllocated(containers);
    Assert.assertEquals("Wrong container allocation count", 1, master.getAllocatedContainers());
    Mockito.verifyZeroInteractions(mockClient);
    Assert.assertEquals("Incorrect number of threads launched", 1, master.threadsLaunched);
    Assert.assertEquals("Incorrect YARN Shell IDs", Arrays.asList("1"), master.yarnShellIds);
    // now send 3 extra containers
    containers.clear();
    ContainerId id2 = BuilderUtils.newContainerId(1, 1, 1, 2);
    containers.add(generateContainer(id2));
    ContainerId id3 = BuilderUtils.newContainerId(1, 1, 1, 3);
    containers.add(generateContainer(id3));
    ContainerId id4 = BuilderUtils.newContainerId(1, 1, 1, 4);
    containers.add(generateContainer(id4));
    handler.onContainersAllocated(containers);
    Assert.assertEquals("Wrong final container allocation count", 4, master.getAllocatedContainers());
    Assert.assertEquals("Incorrect number of threads launched", 4, master.threadsLaunched);
    Assert.assertEquals("Incorrect YARN Shell IDs", Arrays.asList("1", "2", "3", "4"), master.yarnShellIds);
    // make sure we handle completion events correctly
    List<ContainerStatus> status = new ArrayList<>();
    status.add(generateContainerStatus(id1, ContainerExitStatus.SUCCESS));
    status.add(generateContainerStatus(id2, ContainerExitStatus.SUCCESS));
    status.add(generateContainerStatus(id3, ContainerExitStatus.ABORTED));
    status.add(generateContainerStatus(id4, ContainerExitStatus.ABORTED));
    handler.onContainersCompleted(status);
    Assert.assertEquals("Unexpected number of completed containers", targetContainers, master.getNumCompletedContainers());
    Assert.assertTrue("Master didn't finish containers as expected", master.getDone());
    // test for events from containers we know nothing about
    // these events should be ignored
    status = new ArrayList<>();
    ContainerId id5 = BuilderUtils.newContainerId(1, 1, 1, 5);
    status.add(generateContainerStatus(id5, ContainerExitStatus.ABORTED));
    Assert.assertEquals("Unexpected number of completed containers", targetContainers, master.getNumCompletedContainers());
    Assert.assertTrue("Master didn't finish containers as expected", master.getDone());
    status.add(generateContainerStatus(id5, ContainerExitStatus.SUCCESS));
    Assert.assertEquals("Unexpected number of completed containers", targetContainers, master.getNumCompletedContainers());
    Assert.assertTrue("Master didn't finish containers as expected", master.getDone());
}