Examples with Priority org.apache.hadoop.yarn.api.records.Priority used on opensource projects

Example 76 with Priority

use of org.apache.hadoop.yarn.api.records.Priority in project tez by apache.

the class TestContainerReuse method testAssignmentOnShutdown.

@Test(timeout = 10000l)
public void testAssignmentOnShutdown() throws IOException, InterruptedException, ExecutionException {
    LOG.info("Test testAssignmentOnShutdown");
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);
    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.SUCCEEDED).when(appContext).getAMState();
    doReturn(true).when(appContext).isAMInCompletionState();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();
    TaskSchedulerManager taskSchedulerManagerReal = new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(tezConf));
    TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
    taskSchedulerManager.init(tezConf);
    taskSchedulerManager.start();
    TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext) ((TaskSchedulerManagerForTest) taskSchedulerManager).getSpyTaskScheduler();
    TaskSchedulerContextDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;
    Resource resource1 = Resource.newInstance(1024, 1);
    String[] host1 = { "host1" };
    String[] racks = { "/default-rack" };
    Priority priority1 = Priority.newInstance(1);
    TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);
    // Vertex 1, Task 1, Attempt 1, host1
    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent1 = createLaunchRequestEvent(taID11, ta11, resource1, host1, racks, priority1);
    taskSchedulerManager.handleEvent(lrEvent1);
    Container container1 = createContainer(1, "host1", resource1, priority1);
    // One container allocated.
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(taskSchedulerManager, times(0)).taskAllocated(eq(0), eq(ta11), any(Object.class), eq(container1));
    taskScheduler.shutdown();
    taskSchedulerManager.close();
}

Example 77 with Priority

use of org.apache.hadoop.yarn.api.records.Priority in project tez by apache.

the class TestContainerReuse method testDifferentResourceContainerReuse.

@Test(timeout = 5000)
public void testDifferentResourceContainerReuse() throws Exception {
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);
    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();
    TaskSchedulerManager taskSchedulerManagerReal = new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(tezConf));
    TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
    taskSchedulerManager.init(tezConf);
    taskSchedulerManager.start();
    TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext) ((TaskSchedulerManagerForTest) taskSchedulerManager).getSpyTaskScheduler();
    TaskSchedulerContextDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;
    Resource resource1 = Resource.newInstance(1024, 1);
    Resource resource2 = Resource.newInstance(2048, 1);
    String[] host1 = { "host1" };
    String[] host2 = { "host2" };
    String[] racks = { "/default-rack" };
    Priority priority1 = Priority.newInstance(1);
    TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);
    // Vertex 1, Task 1, Attempt 1, host1
    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent1 = createLaunchRequestEvent(taID11, ta11, resource1, host1, racks, priority1);
    // Vertex 1, Task 2, Attempt 1, host1
    TezTaskAttemptID taID12 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 2), 1);
    TaskAttempt ta12 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent2 = createLaunchRequestEvent(taID12, ta12, resource1, host1, racks, priority1);
    // Vertex 1, Task 3, Attempt 1, host2
    TezTaskAttemptID taID13 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 3), 1);
    TaskAttempt ta13 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent3 = createLaunchRequestEvent(taID13, ta13, resource2, host2, racks, priority1);
    // Vertex 1, Task 4, Attempt 1, host2
    TezTaskAttemptID taID14 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 4), 1);
    TaskAttempt ta14 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent4 = createLaunchRequestEvent(taID14, ta14, resource2, host2, racks, priority1);
    taskSchedulerManager.handleEvent(lrEvent1);
    taskSchedulerManager.handleEvent(lrEvent2);
    taskSchedulerManager.handleEvent(lrEvent3);
    taskSchedulerManager.handleEvent(lrEvent4);
    Container container1 = createContainer(1, "host1", resource1, priority1);
    Container container2 = createContainer(2, "host2", resource2, priority1);
    // One container allocated, should start ta11
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta11), any(Object.class), eq(container1));
    // Second container allocated, should start ta13
    taskScheduler.onContainersAllocated(Collections.singletonList(container2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta13), any(Object.class), eq(container2));
    // ta11 finished, should start ta12
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta11, true, null, null);
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta12), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();
    // ta13 finished, should start ta14
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta13, container2.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta13, true, null, null);
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta14), any(Object.class), eq(container2));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container2.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();
    // ta12 finished no pending requests, should release container1
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta12, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta12, true, null, null);
    verify(rmClient).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();
    // ta14 finished no pending requests, should release container2.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta14, container2.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta14, true, null, null);
    verify(rmClient).releaseAssignedContainer(eq(container2.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();
    taskScheduler.shutdown();
    taskSchedulerManager.close();
}

Example 78 with Priority

use of org.apache.hadoop.yarn.api.records.Priority in project tez by apache.

the class TestLocalTaskScheduler method maxTasksAllocationsCannotBeExceeded.

@Test(timeout = 5000)
public void maxTasksAllocationsCannotBeExceeded() {
    final int MAX_TASKS = 4;
    TezConfiguration tezConf = new TezConfiguration();
    tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS);
    ApplicationId appId = ApplicationId.newInstance(2000, 1);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
    TaskSchedulerContext mockContext = TestTaskSchedulerHelpers.setupMockTaskSchedulerContext("", 0, "", true, appAttemptId, 1000l, null, new Configuration());
    LocalContainerFactory containerFactory = new LocalContainerFactory(appAttemptId, 1000);
    HashMap<Object, AllocatedTask> taskAllocations = new LinkedHashMap<>();
    LinkedBlockingQueue<SchedulerRequest> clientRequestQueue = new LinkedBlockingQueue<>();
    // Object under test
    AsyncDelegateRequestHandler requestHandler = new AsyncDelegateRequestHandler(clientRequestQueue, containerFactory, taskAllocations, mockContext, tezConf);
    // Allocate up to max tasks
    for (int i = 0; i < MAX_TASKS; i++) {
        Priority priority = Priority.newInstance(20);
        requestHandler.addAllocateTaskRequest(new Long(i), null, priority, null);
        requestHandler.dispatchRequest();
        requestHandler.allocateTask();
    }
    // Only MAX_TASKS number of tasks should have been allocated
    Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, taskAllocations.size());
    Assert.assertTrue("Another allocation should not fit", !requestHandler.shouldProcess());
    // Deallocate down to zero
    for (int i = 0; i < MAX_TASKS; i++) {
        requestHandler.addDeallocateTaskRequest(new Long(i));
        requestHandler.dispatchRequest();
    }
    // All allocated tasks should have been removed
    Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size());
}

Example 79 with Priority

use of org.apache.hadoop.yarn.api.records.Priority in project tez by apache.

the class TestDagAwareYarnTaskScheduler method testReuseVertexDescendants.

@Test(timeout = 30000)
public void testReuseVertexDescendants() 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);
    // vertex 0 and vertex 2 are root vertices and vertex 1 is a child of vertex 0
    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(BitSet.valueOf(new long[] { 0x2 }));
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(new BitSet());
    when(mockDagInfo.getVertexDescendants(2)).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());
    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    Priority priorityv2 = Priority.newInstance(3);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t0.task)).thenReturn(0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host3");
    when(mockApp.getVertexIndexForTask(taskv0t1.task)).thenReturn(0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv1t0.task)).thenReturn(1);
    MockTaskInfo taskv2t0 = new MockTaskInfo("taskv2t0", priorityv2, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv2t0.task)).thenReturn(2);
    MockTaskInfo taskv2t1 = new MockTaskInfo("taskv2t1", priorityv2, "host3");
    when(mockApp.getVertexIndexForTask(taskv2t1.task)).thenReturn(2);
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    TaskRequest reqv2t0 = taskRequestCaptor.scheduleTask(taskv2t0);
    taskRequestCaptor.scheduleTask(taskv2t1);
    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 v1t0 is skipped and v2t0 assigned instead
    // since host locality beats rack locality for unblocked vertex v2 and
    // v1 is blocked by pending v0 request
    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(taskv2t0.task, taskv2t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv2t0);
    // finish v2t0 successfully, verify v0t1 is assigned since it is higher
    // priority than v2
    assertTrue(scheduler.deallocateTask(taskv2t0.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);
    // finish v2t0 successfully, verify v1t0 is assigned since it is now unblocked
    assertTrue(scheduler.deallocateTask(taskv0t1.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 v2t1
    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();
}

Example 80 with Priority

use of org.apache.hadoop.yarn.api.records.Priority in project tez by apache.

the class TestDagAwareYarnTaskScheduler method testSessionContainers.

@Test(timeout = 30000)
public void testSessionContainers() 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);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000);
    conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5);
    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);
    when(mockApp.isSession()).thenReturn(true);
    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());
    final String rack1 = "/r1";
    final String rack2 = "/r2";
    final String rack3 = "/r3";
    final String node1Rack1 = "n1r1";
    final String node2Rack1 = "n2r1";
    final String node1Rack2 = "n1r2";
    final String node2Rack2 = "n2r2";
    final String node1Rack3 = "n1r3";
    MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);
    MockDNSToSwitchMapping.addRackMapping(node2Rack2, rack2);
    MockDNSToSwitchMapping.addRackMapping(node1Rack3, rack3);
    Priority priorityv0 = Priority.newInstance(1);
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, node2Rack1, rack1);
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, node1Rack1, rack1);
    MockTaskInfo taskv0t3 = new MockTaskInfo("taskv0t3", priorityv0, node2Rack1, rack1);
    MockTaskInfo taskv0t4 = new MockTaskInfo("taskv0t4", priorityv0, node1Rack2, rack2);
    MockTaskInfo taskv0t5 = new MockTaskInfo("taskv0t5", priorityv0, node2Rack2, rack2);
    MockTaskInfo taskv0t6 = new MockTaskInfo("taskv0t6", priorityv0, node1Rack3, rack3);
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv0t3 = taskRequestCaptor.scheduleTask(taskv0t3);
    TaskRequest reqv0t4 = taskRequestCaptor.scheduleTask(taskv0t4);
    TaskRequest reqv0t5 = taskRequestCaptor.scheduleTask(taskv0t5);
    TaskRequest reqv0t6 = taskRequestCaptor.scheduleTask(taskv0t6);
    List<Container> containers = new ArrayList<>();
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    NodeId n1r1 = NodeId.newInstance(node1Rack1, 1);
    Container container1 = Container.newInstance(cid1, n1r1, null, taskv0t0.capability, priorityv0, null);
    containers.add(container1);
    ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
    NodeId n2r1 = NodeId.newInstance(node2Rack1, 1);
    Container container2 = Container.newInstance(cid2, n2r1, null, taskv0t1.capability, priorityv0, null);
    containers.add(container2);
    ContainerId cid3 = ContainerId.newContainerId(attemptId, 3);
    Container container3 = Container.newInstance(cid3, n1r1, null, taskv0t2.capability, priorityv0, null);
    containers.add(container3);
    ContainerId cid4 = ContainerId.newContainerId(attemptId, 4);
    Container container4 = Container.newInstance(cid4, n2r1, null, taskv0t3.capability, priorityv0, null);
    containers.add(container4);
    ContainerId cid5 = ContainerId.newContainerId(attemptId, 5);
    NodeId n1r2 = NodeId.newInstance(node1Rack2, 1);
    Container container5 = Container.newInstance(cid5, n1r2, null, taskv0t4.capability, priorityv0, null);
    containers.add(container5);
    ContainerId cid6 = ContainerId.newContainerId(attemptId, 6);
    NodeId n2r2 = NodeId.newInstance(node2Rack2, 1);
    Container container6 = Container.newInstance(cid6, n2r2, null, taskv0t5.capability, priorityv0, null);
    containers.add(container6);
    ContainerId cid7 = ContainerId.newContainerId(attemptId, 7);
    NodeId n1r3 = NodeId.newInstance(node1Rack3, 1);
    Container container7 = Container.newInstance(cid7, n1r3, null, taskv0t6.capability, priorityv0, null);
    containers.add(container7);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container2);
    verify(mockRMClient).removeContainerRequest(reqv0t1);
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container3);
    verify(mockRMClient).removeContainerRequest(reqv0t2);
    verify(mockApp).taskAllocated(taskv0t3.task, taskv0t3.cookie, container4);
    verify(mockRMClient).removeContainerRequest(reqv0t3);
    verify(mockApp).taskAllocated(taskv0t4.task, taskv0t4.cookie, container5);
    verify(mockRMClient).removeContainerRequest(reqv0t4);
    verify(mockApp).taskAllocated(taskv0t5.task, taskv0t5.cookie, container6);
    verify(mockRMClient).removeContainerRequest(reqv0t5);
    verify(mockApp).taskAllocated(taskv0t6.task, taskv0t6.cookie, container7);
    verify(mockRMClient).removeContainerRequest(reqv0t6);
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t2.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t3.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t4.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t5.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t6.task, true, null, null));
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    // verify only two of the containers were released after idle expiration
    // and the rest were spread across the nodes and racks
    clock.incrementTime(5000);
    drainableAppCallback.drain();
    verify(mockApp, times(2)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(2)).releaseAssignedContainer(any(ContainerId.class));
    Set<String> hosts = new HashSet<>();
    Set<String> racks = new HashSet<>();
    for (HeldContainer hc : scheduler.getSessionContainers()) {
        hosts.add(hc.getHost());
        racks.add(hc.getRack());
    }
    assertEquals(5, hosts.size());
    assertEquals(3, racks.size());
    assertTrue(hosts.contains(node1Rack1));
    assertTrue(hosts.contains(node2Rack1));
    assertTrue(hosts.contains(node1Rack2));
    assertTrue(hosts.contains(node2Rack2));
    assertTrue(hosts.contains(node1Rack3));
    assertTrue(racks.contains(rack1));
    assertTrue(racks.contains(rack2));
    assertTrue(racks.contains(rack3));
    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();
}