Example usage for org.apache.hadoop.yarn.api.records NodeId newInstance

List of usage examples for org.apache.hadoop.yarn.api.records NodeId newInstance

Introduction

In this page you can find the example usage for org.apache.hadoop.yarn.api.records NodeId newInstance.

Prototype

@Public
    @Stable
    public static NodeId newInstance(String host, int port) 

Source Link

Usage

From source file:org.apache.tez.dag.app.rm.node.TestAMNodeMap.java

License:Apache License

@Test
@SuppressWarnings({ "resource", "rawtypes" })
public void testHealthUpdateKnownNode() {

    DrainDispatcher dispatcher = new DrainDispatcher();
    dispatcher.init(new Configuration());
    dispatcher.start();//ww w.j a va  2s .c o m
    EventHandler eventHandler = dispatcher.getEventHandler();

    AppContext appContext = mock(AppContext.class);

    AMNodeMap amNodeMap = new AMNodeMap(eventHandler, appContext);

    NodeId nodeId = NodeId.newInstance("host1", 2342);
    amNodeMap.nodeSeen(nodeId);

    NodeReport nodeReport = generateNodeReport(nodeId, NodeState.UNHEALTHY);
    amNodeMap.handle(new AMNodeEventStateChanged(nodeReport));
    dispatcher.await();
    assertEquals(AMNodeState.UNHEALTHY, amNodeMap.get(nodeId).getState());
    dispatcher.stop();
}

From source file:org.apache.tez.dag.app.rm.node.TestAMNodeMap.java

License:Apache License

@Test
@SuppressWarnings({ "resource", "rawtypes" })
public void testHealthUpdateUnknownNode() {
    DrainDispatcher dispatcher = new DrainDispatcher();
    EventHandler eventHandler = dispatcher.getEventHandler();

    AppContext appContext = mock(AppContext.class);

    AMNodeMap amNodeMap = new AMNodeMap(eventHandler, appContext);

    NodeId nodeId = NodeId.newInstance("unknownhost", 2342);

    NodeReport nodeReport = generateNodeReport(nodeId, NodeState.UNHEALTHY);
    amNodeMap.handle(new AMNodeEventStateChanged(nodeReport));

    // No exceptions - the status update was ignored. Not bothering to capture
    // the log message for verification.
    dispatcher.stop();//from   w w  w.  ja v a2  s.  c o m
}

From source file:org.apache.tez.dag.app.rm.node.TestAMNodeTracker.java

License:Apache License

@Test(timeout = 5000)
public void testHealthUpdateKnownNode() {
    AppContext appContext = mock(AppContext.class);

    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    amNodeTracker.init(new Configuration(false));
    amNodeTracker.start();// ww w.j  a v a  2s  .com

    NodeId nodeId = NodeId.newInstance("host1", 2342);
    amNodeTracker.nodeSeen(nodeId);

    NodeReport nodeReport = generateNodeReport(nodeId, NodeState.UNHEALTHY);
    amNodeTracker.handle(new AMNodeEventStateChanged(nodeReport));
    dispatcher.await();
    assertEquals(AMNodeState.UNHEALTHY, amNodeTracker.get(nodeId).getState());
    amNodeTracker.stop();
}

From source file:org.apache.tez.dag.app.rm.node.TestAMNodeTracker.java

License:Apache License

@Test(timeout = 5000)
public void testHealthUpdateUnknownNode() {
    AppContext appContext = mock(AppContext.class);

    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    amNodeTracker.init(new Configuration(false));
    amNodeTracker.start();//from  w w  w.java  2  s .co  m

    NodeId nodeId = NodeId.newInstance("unknownhost", 2342);

    NodeReport nodeReport = generateNodeReport(nodeId, NodeState.UNHEALTHY);
    amNodeTracker.handle(new AMNodeEventStateChanged(nodeReport));
    dispatcher.await();

    amNodeTracker.stop();
    // No exceptions - the status update was ignored. Not bothering to capture
    // the log message for verification.
}

From source file:org.apache.tez.dag.app.rm.node.TestAMNodeTracker.java

License:Apache License

@Test(timeout = 5000)
public void testSingleNodeNotBlacklisted() {
    AppContext appContext = mock(AppContext.class);
    Configuration conf = new Configuration(false);
    conf.setInt(TezConfiguration.TEZ_AM_MAX_TASK_FAILURES_PER_NODE, 2);
    conf.setBoolean(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_IGNORE_THRESHOLD, 33);

    TestEventHandler handler = new TestEventHandler();
    AMNodeTracker amNodeTracker = new AMNodeTracker(handler, appContext);
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    AMContainerMap amContainerMap = mock(AMContainerMap.class);
    TaskSchedulerEventHandler taskSchedulerEventHandler = mock(TaskSchedulerEventHandler.class);
    dispatcher.register(AMNodeEventType.class, amNodeTracker);
    dispatcher.register(AMContainerEventType.class, amContainerMap);
    dispatcher.register(AMSchedulerEventType.class, taskSchedulerEventHandler);
    amNodeTracker.init(conf);//w ww .  j  a v  a2  s  . c  om
    amNodeTracker.start();

    amNodeTracker.handle(new AMNodeEventNodeCountUpdated(1));
    NodeId nodeId = NodeId.newInstance("host1", 1234);
    amNodeTracker.nodeSeen(nodeId);

    AMNodeImpl node = (AMNodeImpl) amNodeTracker.get(nodeId);

    ContainerId cId1 = mock(ContainerId.class);
    ContainerId cId2 = mock(ContainerId.class);

    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId, cId1));
    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId, cId2));

    TezTaskAttemptID ta1 = mock(TezTaskAttemptID.class);
    TezTaskAttemptID ta2 = mock(TezTaskAttemptID.class);

    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId, cId1, ta1, true));
    dispatcher.await();
    assertEquals(1, node.numFailedTAs);
    assertEquals(AMNodeState.ACTIVE, node.getState());

    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId, cId2, ta2, true));
    dispatcher.await();
    assertEquals(2, node.numFailedTAs);
    assertEquals(1, handler.events.size());
    assertEquals(AMNodeEventType.N_IGNORE_BLACKLISTING_ENABLED, handler.events.get(0).getType());
    assertEquals(AMNodeState.FORCED_ACTIVE, node.getState());
}

From source file:org.apache.tez.dag.app.rm.node.TestAMNodeTracker.java

License:Apache License

@Test(timeout = 10000)
public void testNodeSelfBlacklist() {
    AppContext appContext = mock(AppContext.class);
    Configuration conf = new Configuration(false);
    conf.setInt(TezConfiguration.TEZ_AM_MAX_TASK_FAILURES_PER_NODE, 2);
    TestEventHandler handler = new TestEventHandler();
    AMNodeTracker amNodeTracker = new AMNodeTracker(handler, appContext);
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    AMContainerMap amContainerMap = mock(AMContainerMap.class);
    TaskSchedulerEventHandler taskSchedulerEventHandler = mock(TaskSchedulerEventHandler.class);
    dispatcher.register(AMNodeEventType.class, amNodeTracker);
    dispatcher.register(AMContainerEventType.class, amContainerMap);
    dispatcher.register(AMSchedulerEventType.class, taskSchedulerEventHandler);
    amNodeTracker.init(conf);/*from w  w w .  j  a v  a  2s  .  c  o m*/
    amNodeTracker.start();

    amNodeTracker.handle(new AMNodeEventNodeCountUpdated(4));
    NodeId nodeId = NodeId.newInstance("host1", 1234);
    NodeId nodeId2 = NodeId.newInstance("host2", 1234);
    NodeId nodeId3 = NodeId.newInstance("host3", 1234);
    NodeId nodeId4 = NodeId.newInstance("host4", 1234);
    amNodeTracker.nodeSeen(nodeId);
    amNodeTracker.nodeSeen(nodeId2);
    amNodeTracker.nodeSeen(nodeId3);
    amNodeTracker.nodeSeen(nodeId4);
    AMNodeImpl node = (AMNodeImpl) amNodeTracker.get(nodeId);

    ContainerId cId1 = mock(ContainerId.class);
    ContainerId cId2 = mock(ContainerId.class);
    ContainerId cId3 = mock(ContainerId.class);

    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId, cId1));
    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId, cId2));
    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId, cId3));
    assertEquals(3, node.containers.size());

    TezTaskAttemptID ta1 = mock(TezTaskAttemptID.class);
    TezTaskAttemptID ta2 = mock(TezTaskAttemptID.class);
    TezTaskAttemptID ta3 = mock(TezTaskAttemptID.class);

    amNodeTracker.handle(new AMNodeEventTaskAttemptSucceeded(nodeId, cId1, ta1));
    assertEquals(1, node.numSuccessfulTAs);

    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId, cId2, ta2, true));
    assertEquals(1, node.numSuccessfulTAs);
    assertEquals(1, node.numFailedTAs);
    assertEquals(AMNodeState.ACTIVE, node.getState());
    // duplicate should not affect anything
    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId, cId2, ta2, true));
    assertEquals(1, node.numSuccessfulTAs);
    assertEquals(1, node.numFailedTAs);
    assertEquals(AMNodeState.ACTIVE, node.getState());

    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId, cId3, ta3, true));
    dispatcher.await();
    assertEquals(1, node.numSuccessfulTAs);
    assertEquals(2, node.numFailedTAs);
    assertEquals(AMNodeState.BLACKLISTED, node.getState());

    assertEquals(4, handler.events.size());
    assertEquals(AMContainerEventType.C_NODE_FAILED, handler.events.get(0).getType());
    assertEquals(cId1, ((AMContainerEventNodeFailed) handler.events.get(0)).getContainerId());
    assertEquals(AMContainerEventType.C_NODE_FAILED, handler.events.get(1).getType());
    assertEquals(cId2, ((AMContainerEventNodeFailed) handler.events.get(1)).getContainerId());
    assertEquals(AMContainerEventType.C_NODE_FAILED, handler.events.get(2).getType());
    assertEquals(cId3, ((AMContainerEventNodeFailed) handler.events.get(2)).getContainerId());
    assertEquals(AMSchedulerEventType.S_NODE_BLACKLISTED, handler.events.get(3).getType());
    assertEquals(node.getNodeId(), ((AMSchedulerEventNodeBlacklistUpdate) handler.events.get(3)).getNodeId());

    // Trigger one more node failure, which should cause BLACKLISTING to be disabled
    ContainerId cId4 = mock(ContainerId.class);
    ContainerId cId5 = mock(ContainerId.class);
    TezTaskAttemptID ta4 = mock(TezTaskAttemptID.class);
    TezTaskAttemptID ta5 = mock(TezTaskAttemptID.class);
    AMNodeImpl node2 = (AMNodeImpl) amNodeTracker.get(nodeId2);
    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId2, cId4));
    amNodeTracker.handle(new AMNodeEventContainerAllocated(nodeId2, cId5));

    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId2, cId4, ta4, true));
    assertEquals(1, node2.numFailedTAs);
    assertEquals(AMNodeState.ACTIVE, node2.getState());

    handler.events.clear();
    amNodeTracker.handle(new AMNodeEventTaskAttemptEnded(nodeId2, cId5, ta5, true));
    dispatcher.await();
    assertEquals(2, node2.numFailedTAs);
    assertEquals(AMNodeState.FORCED_ACTIVE, node2.getState());
    AMNodeImpl node3 = (AMNodeImpl) amNodeTracker.get(nodeId3);
    assertEquals(AMNodeState.FORCED_ACTIVE, node3.getState());
    assertEquals(5, handler.events.size());

    // Blacklisting Disabled, the node causing this will not be blacklisted. The single node that
    // was blacklisted will be unblacklisted.
    int numIgnoreBlacklistingEnabledEvents = 0;
    int numUnblacklistedEvents = 0;
    for (Event event : handler.events) {
        if (event.getType() == AMNodeEventType.N_IGNORE_BLACKLISTING_ENABLED) {
            numIgnoreBlacklistingEnabledEvents++;
        } else if (event.getType() == AMSchedulerEventType.S_NODE_UNBLACKLISTED) {
            numUnblacklistedEvents++;
        } else {
            fail("Unexpected event of type: " + event.getType());
        }
    }
    assertEquals(4, numIgnoreBlacklistingEnabledEvents);
    assertEquals(1, numUnblacklistedEvents);

    // drain all previous events
    dispatcher.await();

    // Increase the number of nodes. BLACKLISTING should be re-enabled.
    // Node 1 and Node 2 should go into BLACKLISTED state
    handler.events.clear();
    amNodeTracker.handle(new AMNodeEventNodeCountUpdated(8));
    dispatcher.await();
    LOG.info(("Completed waiting for dispatcher to process all pending events"));
    assertEquals(AMNodeState.BLACKLISTED, node.getState());
    assertEquals(AMNodeState.BLACKLISTED, node2.getState());
    assertEquals(AMNodeState.ACTIVE, node3.getState());
    assertEquals(8, handler.events.size());

    int index = 0;
    int numIgnoreBlacklistingDisabledEvents = 0;
    int numBlacklistedEvents = 0;
    int numNodeFailedEvents = 0;
    for (Event event : handler.events) {
        LOG.info("Logging event: index:" + index++ + " type: " + event.getType());
        if (event.getType() == AMNodeEventType.N_IGNORE_BLACKLISTING_DISABLED) {
            numIgnoreBlacklistingDisabledEvents++;
        } else if (event.getType() == AMSchedulerEventType.S_NODE_BLACKLISTED) {
            numBlacklistedEvents++;
        } else if (event.getType() == AMContainerEventType.C_NODE_FAILED) {
            numNodeFailedEvents++;
            // Node2 is now blacklisted so the container's will be informed
            assertTrue(((AMContainerEventNodeFailed) event).getContainerId() == cId4
                    || ((AMContainerEventNodeFailed) event).getContainerId() == cId5);
        } else {
            fail("Unexpected event of type: " + event.getType());
        }
    }
    assertEquals(4, numIgnoreBlacklistingDisabledEvents);
    assertEquals(2, numBlacklistedEvents);
    assertEquals(2, numNodeFailedEvents);

    amNodeTracker.stop();
}

From source file:org.apache.tez.dag.app.rm.TestContainerReuse.java

License:Apache License

private Container createContainer(int id, String host, Resource resource, Priority priority) {
    ContainerId containerID = ContainerId
            .newInstance(ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1), id);
    NodeId nodeID = NodeId.newInstance(host, 0);
    Container container = Container.newInstance(containerID, nodeID, host + ":0", resource, priority, null);
    return container;
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@SuppressWarnings({ "unchecked" })
@Test(timeout = 30000)/*from   ww  w.  j  av a 2s.c  o m*/
public void testNoReuse() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Object mockTask1 = new MockTask("task1");
    Object mockCookie1 = new Object();
    Resource mockCapability = Resources.createResource(1024, 1);
    String[] hosts = { "host1", "host5" };
    String[] racks = { "/default-rack", "/default-rack" };
    Priority mockPriority = Priority.newInstance(1);
    ArgumentCaptor<TaskRequest> requestCaptor = ArgumentCaptor.forClass(TaskRequest.class);
    // allocate task
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).addContainerRequest(any(TaskRequest.class));

    // returned from task requests before allocation happens
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
    verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(1)).removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());

    // deallocating unknown task
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
    verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(1)).removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());

    // allocate tasks
    Object mockTask2 = new MockTask("task2");
    Object mockCookie2 = new Object();
    Object mockTask3 = new MockTask("task3");
    Object mockCookie3 = new Object();
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture());
    TaskRequest request1 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask2, mockCapability, hosts, racks, mockPriority, null, mockCookie2);
    drainableAppCallback.drain();
    verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
    TaskRequest request2 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask3, mockCapability, hosts, racks, mockPriority, null, mockCookie3);
    drainableAppCallback.drain();
    verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture());
    TaskRequest request3 = requestCaptor.getValue();

    NodeId host1 = NodeId.newInstance("host1", 1);
    NodeId host2 = NodeId.newInstance("host2", 2);
    NodeId host3 = NodeId.newInstance("host3", 3);
    NodeId host4 = NodeId.newInstance("host4", 4);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId mockCId1 = ContainerId.newContainerId(attemptId, 1);
    Container mockContainer1 = Container.newInstance(mockCId1, host1, null, mockCapability, mockPriority, null);
    ContainerId mockCId2 = ContainerId.newContainerId(attemptId, 2);
    Container mockContainer2 = Container.newInstance(mockCId2, host2, null, mockCapability, mockPriority, null);
    ContainerId mockCId3 = ContainerId.newContainerId(attemptId, 3);
    Container mockContainer3 = Container.newInstance(mockCId3, host3, null, mockCapability, mockPriority, null);
    ContainerId mockCId4 = ContainerId.newContainerId(attemptId, 4);
    Container mockContainer4 = Container.newInstance(mockCId4, host4, null, mockCapability, mockPriority, null);
    List<Container> containers = new ArrayList<>();
    containers.add(mockContainer1);
    containers.add(mockContainer2);
    containers.add(mockContainer3);
    containers.add(mockContainer4);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // first container allocated
    verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1);
    verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2);
    verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3);
    // no other allocations returned
    verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
    verify(mockRMClient).removeContainerRequest(request1);
    verify(mockRMClient).removeContainerRequest(request2);
    verify(mockRMClient).removeContainerRequest(request3);
    // verify unwanted container released
    verify(mockRMClient).releaseAssignedContainer(mockCId4);

    // deallocate allocated task
    assertTrue(scheduler.deallocateTask(mockTask1, true, null, null));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(mockCId1);
    verify(mockRMClient).releaseAssignedContainer(mockCId1);
    // deallocate allocated container
    assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2));
    drainableAppCallback.drain();
    verify(mockRMClient).releaseAssignedContainer(mockCId2);
    verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());

    List<ContainerStatus> statuses = new ArrayList<>();
    ContainerStatus mockStatus1 = mock(ContainerStatus.class);
    when(mockStatus1.getContainerId()).thenReturn(mockCId1);
    statuses.add(mockStatus1);
    ContainerStatus mockStatus2 = mock(ContainerStatus.class);
    when(mockStatus2.getContainerId()).thenReturn(mockCId2);
    statuses.add(mockStatus2);
    ContainerStatus mockStatus3 = mock(ContainerStatus.class);
    when(mockStatus3.getContainerId()).thenReturn(mockCId3);
    statuses.add(mockStatus3);
    ContainerStatus mockStatus4 = mock(ContainerStatus.class);
    when(mockStatus4.getContainerId()).thenReturn(mockCId4);
    statuses.add(mockStatus4);

    scheduler.onContainersCompleted(statuses);
    drainableAppCallback.drain();
    // released container status returned
    verify(mockApp).containerCompleted(mockTask1, mockStatus1);
    verify(mockApp).containerCompleted(mockTask2, mockStatus2);
    // currently allocated container status returned and not released
    verify(mockApp).containerCompleted(mockTask3, mockStatus3);
    // no other statuses returned
    verify(mockApp, times(3)).containerCompleted(any(), (ContainerStatus) any());
    verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());

    // verify blacklisting
    verify(mockRMClient, times(0)).updateBlacklist(anyListOf(String.class), anyListOf(String.class));
    String badHost = "host6";
    NodeId badNodeId = NodeId.newInstance(badHost, 1);
    scheduler.blacklistNode(badNodeId);
    List<String> badNodeList = Collections.singletonList(badHost);
    verify(mockRMClient, times(1)).updateBlacklist(eq(badNodeList), isNull(List.class));
    Object mockTask4 = new MockTask("task4");
    Object mockCookie4 = new Object();
    scheduler.allocateTask(mockTask4, mockCapability, null, null, mockPriority, null, mockCookie4);
    drainableAppCallback.drain();
    verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture());
    ContainerId mockCId5 = ContainerId.newContainerId(attemptId, 5);
    Container mockContainer5 = Container.newInstance(mockCId5, badNodeId, null, mockCapability, mockPriority,
            null);
    containers.clear();
    containers.add(mockContainer5);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // no new allocation
    verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
    // verify blacklisted container released
    verify(mockRMClient).releaseAssignedContainer(mockCId5);
    verify(mockRMClient, times(4)).releaseAssignedContainer((ContainerId) any());
    // verify request added back
    verify(mockRMClient, times(6)).addContainerRequest(requestCaptor.capture());
    NodeId host6 = NodeId.newInstance("host6", 6);
    ContainerId mockCId6 = ContainerId.newContainerId(attemptId, 6);
    Container mockContainer6 = Container.newInstance(mockCId6, host6, null, mockCapability, mockPriority, null);
    containers.clear();
    containers.add(mockContainer6);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // new allocation
    verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any());
    verify(mockApp).taskAllocated(mockTask4, mockCookie4, mockContainer6);
    // deallocate allocated task
    assertTrue(scheduler.deallocateTask(mockTask4, true, null, null));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(mockCId6);
    verify(mockRMClient).releaseAssignedContainer(mockCId6);
    verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any());
    // test unblacklist
    scheduler.unblacklistNode(badNodeId);
    verify(mockRMClient, times(1)).updateBlacklist(isNull(List.class), eq(badNodeList));
    assertEquals(0, scheduler.getNumBlacklistedNodes());

    float progress = 0.5f;
    when(mockApp.getProgress()).thenReturn(progress);
    assertEquals(progress, scheduler.getProgress(), 0);

    // check duplicate allocation request
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(7)).addContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(6)).removeContainerRequest(any(TaskRequest.class));
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    // old request removed and new one added
    verify(mockRMClient, times(7)).removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(8)).addContainerRequest(any(TaskRequest.class));
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));

    List<NodeReport> mockUpdatedNodes = mock(List.class);
    scheduler.onNodesUpdated(mockUpdatedNodes);
    drainableAppCallback.drain();
    verify(mockApp).nodesUpdated(mockUpdatedNodes);

    ArgumentCaptor<String> argumentCaptor = ArgumentCaptor.forClass(String.class);
    Exception mockException = new IOException("mockexception");
    scheduler.onError(mockException);
    drainableAppCallback.drain();
    verify(mockApp).reportError(eq(YarnTaskSchedulerServiceError.RESOURCEMANAGER_ERROR),
            argumentCaptor.capture(), any(DagInfo.class));
    assertTrue(argumentCaptor.getValue().contains("mockexception"));

    scheduler.onShutdownRequest();
    drainableAppCallback.drain();
    verify(mockApp).appShutdownRequested();

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testSimpleReuseLocalMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();/*w  w w .  ja v  a2s  . c o  m*/
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host3");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, hostsv0t0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, hostsv0t0);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v0t2 instead is assigned to the container
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t2);

    // finish v0t2 successfully, verify v1t0 is assigned to the same container
    assertTrue(scheduler.deallocateTask(taskv0t2.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // fail v1t0 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv1t0.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testSimpleReuseRackMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();/*from   w  w  w. j  ava  2s  .c o  m*/
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host2");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, "host4", "/somerack");
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, "host5");

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v1t0 assigned instead
    // since host locality is preferred to rack locality and lower priority vertex
    // is not blocked by higher priority vertex
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // finish v1t0 successfully, verify v0t1 is assigned
    assertTrue(scheduler.deallocateTask(taskv1t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // fail v0t1 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv0t1.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}