Examples with TaskSchedulerContext org.apache.tez.serviceplugins.api.TaskSchedulerContext used on opensource projects

Example 16 with TaskSchedulerContext

use of org.apache.tez.serviceplugins.api.TaskSchedulerContext in project tez by apache.

the class TestDagAwareYarnTaskScheduler method testSimpleReuseAnyMatching.

@Test(timeout = 30000)
public void testSimpleReuseAnyMatching() 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, true);
    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());
    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", "/rack4");
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, "host6", "/rack6");
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = 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 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);
    // finish v0t1 successfully, verify v0t2 is assigned
    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(taskv0t2.task, taskv0t2.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t2);
    // fail v0t2 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv0t2.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 17 with TaskSchedulerContext

use of org.apache.tez.serviceplugins.api.TaskSchedulerContext in project tez by apache.

the class TestDagAwareYarnTaskScheduler method testReuseWithAffinity.

@Test(timeout = 30000)
public void testReuseWithAffinity() 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();
    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, hostsv0t0);
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    taskRequestCaptor.scheduleTask(taskv0t1);
    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);
    // add a new request for this container
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0, cid1);
    // finish v0t0 successfully, verify v0t1 is skipped even though it is node-local
    // and v1t0 assigned instead for affinity
    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);
    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 18 with TaskSchedulerContext

use of org.apache.tez.serviceplugins.api.TaskSchedulerContext in project tez by apache.

the class TestLocalTaskSchedulerService method testDeallocationAfterAllocation.

/**
 * TaskAttempt Killed from START_WAIT
 */
@Test(timeout = 5000)
public void testDeallocationAfterAllocation() throws InterruptedException {
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(10000l, 1), 1);
    TaskSchedulerContext mockContext = TestTaskSchedulerHelpers.setupMockTaskSchedulerContext("", 0, "", false, appAttemptId, 10000l, null, new Configuration());
    MockLocalTaskSchedulerSerivce taskSchedulerService = new MockLocalTaskSchedulerSerivce(mockContext);
    taskSchedulerService.initialize();
    taskSchedulerService.start();
    Task task = mock(Task.class);
    taskSchedulerService.allocateTask(task, Resource.newInstance(1024, 1), null, null, Priority.newInstance(1), null, null);
    taskSchedulerService.startRequestHandlerThread();
    MockAsyncDelegateRequestHandler requestHandler = taskSchedulerService.getRequestHandler();
    requestHandler.drainRequest(1);
    taskSchedulerService.deallocateTask(task, false, null, null);
    requestHandler.drainRequest(2);
    assertEquals(1, requestHandler.deallocateCount);
    assertEquals(1, requestHandler.allocateCount);
    taskSchedulerService.shutdown();
}

Example 19 with TaskSchedulerContext

use of org.apache.tez.serviceplugins.api.TaskSchedulerContext in project tez by apache.

the class TestLocalTaskSchedulerService method preemptDescendantsOnly.

@Test
public void preemptDescendantsOnly() {
    final int MAX_TASKS = 2;
    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);
    Long parentTask1 = new Long(1);
    Long parentTask2 = new Long(2);
    Long childTask1 = new Long(3);
    Long grandchildTask1 = new Long(4);
    TaskSchedulerContext mockContext = TestTaskSchedulerHelpers.setupMockTaskSchedulerContext("", 0, "", true, appAttemptId, 1000l, null, tezConf);
    when(mockContext.getVertexIndexForTask(parentTask1)).thenReturn(0);
    when(mockContext.getVertexIndexForTask(parentTask2)).thenReturn(0);
    when(mockContext.getVertexIndexForTask(childTask1)).thenReturn(1);
    when(mockContext.getVertexIndexForTask(grandchildTask1)).thenReturn(2);
    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    BitSet vertex1Descendants = new BitSet();
    vertex1Descendants.set(1);
    vertex1Descendants.set(2);
    BitSet vertex2Descendants = new BitSet();
    vertex2Descendants.set(2);
    BitSet vertex3Descendants = new BitSet();
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(vertex1Descendants);
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(vertex2Descendants);
    when(mockDagInfo.getVertexDescendants(2)).thenReturn(vertex3Descendants);
    when(mockContext.getCurrentDagInfo()).thenReturn(mockDagInfo);
    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);
    Priority priority3 = Priority.newInstance(3);
    Priority priority4 = Priority.newInstance(4);
    Resource resource = Resource.newInstance(1024, 1);
    MockLocalTaskSchedulerSerivce taskSchedulerService = new MockLocalTaskSchedulerSerivce(mockContext);
    // The mock context need to send a deallocate container request to the scheduler service
    Answer<Void> answer = new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) {
            ContainerId containerId = invocation.getArgumentAt(0, ContainerId.class);
            taskSchedulerService.deallocateContainer(containerId);
            return null;
        }
    };
    doAnswer(answer).when(mockContext).preemptContainer(any(ContainerId.class));
    taskSchedulerService.initialize();
    taskSchedulerService.start();
    taskSchedulerService.startRequestHandlerThread();
    MockAsyncDelegateRequestHandler requestHandler = taskSchedulerService.getRequestHandler();
    taskSchedulerService.allocateTask(parentTask1, resource, null, null, priority1, null, null);
    taskSchedulerService.allocateTask(childTask1, resource, null, null, priority3, null, null);
    taskSchedulerService.allocateTask(grandchildTask1, resource, null, null, priority4, null, null);
    requestHandler.drainRequest(3);
    // We should not preempt if we have not reached max task allocations
    Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, requestHandler.allocateCount);
    Assert.assertTrue("Another allocation should not fit", !requestHandler.shouldProcess());
    // Next task allocation should preempt
    taskSchedulerService.allocateTask(parentTask2, Resource.newInstance(1024, 1), null, null, priority2, null, null);
    requestHandler.drainRequest(5);
    // All allocated tasks should have been removed
    Assert.assertEquals("Wrong number of preempted tasks", 1, requestHandler.preemptCount);
}

Example 20 with TaskSchedulerContext

use of org.apache.tez.serviceplugins.api.TaskSchedulerContext in project tez by apache.

the class TestTaskScheduler method testTaskSchedulerInitiateStop.

@Test(timeout = 10000)
public void testTaskSchedulerInitiateStop() throws Exception {
    Configuration conf = new Configuration();
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
    // keep containers held for 10 seconds
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 10000);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 10000);
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL, conf);
    final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
    TezAMRMClientAsync<CookieContainerRequest> mockRMClient = spy(new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));
    TaskSchedulerWithDrainableContext scheduler = new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);
    scheduler.initialize();
    drainableAppCallback.drain();
    scheduler.start();
    drainableAppCallback.drain();
    Object mockTask1 = mock(Object.class);
    when(mockTask1.toString()).thenReturn("task1");
    Object mockCookie1 = mock(Object.class);
    Resource mockCapability = mock(Resource.class);
    String[] hosts = { "host1", "host5" };
    String[] racks = { "/default-rack", "/default-rack" };
    final Priority mockPriority1 = Priority.newInstance(1);
    final Priority mockPriority2 = Priority.newInstance(2);
    final Priority mockPriority3 = Priority.newInstance(3);
    Object mockTask2 = mock(Object.class);
    when(mockTask2.toString()).thenReturn("task2");
    Object mockCookie2 = mock(Object.class);
    Object mockTask3 = mock(Object.class);
    when(mockTask3.toString()).thenReturn("task3");
    Object mockCookie3 = mock(Object.class);
    ArgumentCaptor<CookieContainerRequest> requestCaptor = ArgumentCaptor.forClass(CookieContainerRequest.class);
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority1, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).addContainerRequest(requestCaptor.capture());
    CookieContainerRequest request1 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask2, mockCapability, hosts, racks, mockPriority2, null, mockCookie2);
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture());
    CookieContainerRequest request2 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask3, mockCapability, hosts, racks, mockPriority3, null, mockCookie3);
    drainableAppCallback.drain();
    verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
    CookieContainerRequest request3 = requestCaptor.getValue();
    List<Container> containers = new ArrayList<Container>();
    // sending lower priority container first to make sure its not matched
    Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer1.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer1.getPriority()).thenReturn(mockPriority1);
    when(mockContainer1.toString()).thenReturn("container1");
    ContainerId mockCId1 = mock(ContainerId.class);
    when(mockContainer1.getId()).thenReturn(mockCId1);
    when(mockCId1.toString()).thenReturn("container1");
    containers.add(mockContainer1);
    Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer2.getNodeId().getHost()).thenReturn("host2");
    when(mockContainer2.getPriority()).thenReturn(mockPriority2);
    when(mockContainer2.toString()).thenReturn("container2");
    ContainerId mockCId2 = mock(ContainerId.class);
    when(mockContainer2.getId()).thenReturn(mockCId2);
    when(mockCId2.toString()).thenReturn("container2");
    containers.add(mockContainer2);
    ArrayList<CookieContainerRequest> hostContainers = new ArrayList<CookieContainerRequest>();
    hostContainers.add(request1);
    ArrayList<CookieContainerRequest> rackContainers = new ArrayList<CookieContainerRequest>();
    rackContainers.add(request2);
    ArrayList<CookieContainerRequest> anyContainers = new ArrayList<CookieContainerRequest>();
    anyContainers.add(request3);
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    scheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;
    scheduler.onContainersAllocated(containers);
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    Assert.assertEquals(2, scheduler.heldContainers.size());
    Assert.assertEquals(1, scheduler.taskRequests.size());
    // 2 containers are allocated and their corresponding taskRequests are removed.
    verify(mockRMClient).removeContainerRequest(request1);
    verify(mockRMClient).removeContainerRequest(request2);
    scheduler.initiateStop();
    // verify all the containers are released
    Assert.assertEquals(0, scheduler.heldContainers.size());
    verify(mockRMClient).releaseAssignedContainer(mockCId1);
    verify(mockRMClient).releaseAssignedContainer(mockCId2);
    // verify taskRequests are removed
    Assert.assertEquals(0, scheduler.taskRequests.size());
    verify(mockRMClient).removeContainerRequest(request3);
}