Examples of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableAppCallback

Package org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers

Examples of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableAppCallback

org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableAppCallback

                                                  mock(TezAMRMClientAsync.class);


    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";
    final TaskSchedulerWithDrainableAppCallback scheduler =
      new TaskSchedulerWithDrainableAppCallback(
        mockApp, new PreemptionMatcher(), appHost, appPort,
        appUrl, mockRMClient, mockAppContext);
    TaskSchedulerAppCallbackDrainable drainableAppCallback = scheduler
        .getDrainableAppCallback();


    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
    scheduler.init(conf);


    RegisterApplicationMasterResponse mockRegResponse =
                       mock(RegisterApplicationMasterResponse.class);
    when(
        mockRMClient.registerApplicationMaster(anyString(), anyInt(),
            anyString())).thenReturn(mockRegResponse);


    scheduler.start();
    Resource totalResource = Resource.newInstance(4000, 4);
    when(mockRMClient.getAvailableResources()).thenReturn(totalResource);


    // no preemption
    scheduler.getProgress();
    drainableAppCallback.drain();
    Assert.assertEquals(totalResource, scheduler.getTotalResources());
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId)any());


    // allocate task
    Object mockTask1 = mock(Object.class);
    Object mockTask2 = mock(Object.class);
    Object mockTask3 = mock(Object.class);
    Object mockTask3Wait = mock(Object.class);
    Object mockTask3Retry = mock(Object.class);
    Object mockTask3Kill = mock(Object.class);
    Object obj3 = new Object();
    Priority pri2 = Priority.newInstance(2);
    Priority pri4 = Priority.newInstance(4);
    Priority pri5 = Priority.newInstance(5);
    Priority pri6 = Priority.newInstance(6);


    ArgumentCaptor<CookieContainerRequest> requestCaptor =
        ArgumentCaptor.forClass(CookieContainerRequest.class);
    final ArrayList<CookieContainerRequest> anyContainers =
        new ArrayList<CookieContainerRequest>();


    
    Resource taskAsk = Resource.newInstance(1024, 1);
    scheduler.allocateTask(mockTask1, taskAsk, null,
                           null, pri2, null, null);
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).
        addContainerRequest(requestCaptor.capture());
    anyContainers.add(requestCaptor.getValue());
    scheduler.allocateTask(mockTask3, taskAsk, null,
                           null, pri6, obj3, null);
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).
    addContainerRequest(requestCaptor.capture());
    anyContainers.add(requestCaptor.getValue());
    // later one in the allocation gets killed between the two task3's
    scheduler.allocateTask(mockTask3Kill, taskAsk, null,
                           null, pri6, obj3, null);
    drainableAppCallback.drain();
    verify(mockRMClient, times(3)).
    addContainerRequest(requestCaptor.capture());
    anyContainers.add(requestCaptor.getValue());


    Resource freeResource = Resource.newInstance(500, 0);
    when(mockRMClient.getAvailableResources()).thenReturn(freeResource);
    scheduler.getProgress();
    drainableAppCallback.drain();
    Assert.assertEquals(totalResource, scheduler.getTotalResources());
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId)any());


    final List<ArrayList<CookieContainerRequest>> anyList =
        new LinkedList<ArrayList<CookieContainerRequest>>();
    final List<ArrayList<CookieContainerRequest>> emptyList =
        new LinkedList<ArrayList<CookieContainerRequest>>();


    anyList.add(anyContainers);
    List<Container> containers = new ArrayList<Container>();
    Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer1.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer1.getResource()).thenReturn(taskAsk);
    when(mockContainer1.getPriority()).thenReturn(pri2);
    ContainerId mockCId1 = mock(ContainerId.class);
    when(mockContainer1.getId()).thenReturn(mockCId1);
    containers.add(mockContainer1);
    Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer2.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer2.getResource()).thenReturn(taskAsk);
    when(mockContainer2.getPriority()).thenReturn(pri6);
    ContainerId mockCId2 = mock(ContainerId.class);
    when(mockContainer2.getId()).thenReturn(mockCId2);
    containers.add(mockContainer2);
    Container mockContainer3 = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer3.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer3.getResource()).thenReturn(taskAsk);
    when(mockContainer3.getPriority()).thenReturn(pri6);
    ContainerId mockCId3 = mock(ContainerId.class);
    when(mockContainer3.getId()).thenReturn(mockCId3);
    containers.add(mockContainer3);
    when(
        mockRMClient.getMatchingRequests((Priority) any(), eq("host1"),
            (Resource) any())).thenAnswer(
        new Answer<List<? extends Collection<CookieContainerRequest>>>() {
          @Override
          public List<? extends Collection<CookieContainerRequest>> answer(
              InvocationOnMock invocation) throws Throwable {
            return emptyList;
          }


        });
    // RackResolver by default puts hosts in default-rack
    when(
        mockRMClient.getMatchingRequests((Priority) any(), eq("/default-rack"),
            (Resource) any())).thenAnswer(
        new Answer<List<? extends Collection<CookieContainerRequest>>>() {
          @Override
          public List<? extends Collection<CookieContainerRequest>> answer(
              InvocationOnMock invocation) throws Throwable {
            return emptyList;
          }


        });
    when(
        mockRMClient.getMatchingRequests((Priority) any(),
            eq(ResourceRequest.ANY), (Resource) any())).thenAnswer(
        new Answer<List<? extends Collection<CookieContainerRequest>>>() {
          int calls = 0;
          @Override
          public List<? extends Collection<CookieContainerRequest>> answer(
              InvocationOnMock invocation) throws Throwable {
            if(calls > 0) {
              anyContainers.remove(0);
            }
            calls++;
            return anyList;
          }


        });
    
    Mockito.doAnswer(new Answer() {
      public Object answer(InvocationOnMock invocation) {
          Object[] args = invocation.getArguments();
          ContainerId cId = (ContainerId) args[0];
          scheduler.deallocateContainer(cId);
          return null;
      }})
    .when(mockApp).preemptContainer((ContainerId)any());
    
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    Assert.assertEquals(3, scheduler.taskAllocations.size());
    Assert.assertEquals(3072, scheduler.allocatedResources.getMemory());
    Assert.assertEquals(mockCId1,
        scheduler.taskAllocations.get(mockTask1).getId());
    Assert.assertEquals(mockCId2,
        scheduler.taskAllocations.get(mockTask3).getId());
    Assert.assertEquals(mockCId3,
        scheduler.taskAllocations.get(mockTask3Kill).getId());


    // no preemption
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId)any());


    Object mockTask3WaitCookie = new Object();
    scheduler.allocateTask(mockTask3Wait, taskAsk, null,
                           null, pri6, obj3, mockTask3WaitCookie);
    // no preemption - same pri
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId)any());
    
    Priority pri8 = Priority.newInstance(8);
    Container mockContainer4 = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer4.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer4.getResource()).thenReturn(taskAsk);
    when(mockContainer4.getPriority()).thenReturn(pri8);
    ContainerId mockCId4 = mock(ContainerId.class);
    when(mockContainer4.getId()).thenReturn(mockCId4);
    containers.clear();
    containers.add(mockContainer4);
    
    // Fudge new container being present in delayed allocation list due to race
    HeldContainer heldContainer = new HeldContainer(mockContainer4, -1, -1, null);
    scheduler.delayedContainerManager.delayedContainers.add(heldContainer);
    // no preemption - container assignment attempts < 3
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId)any());
    heldContainer.incrementAssignmentAttempts();
    // no preemption - container assignment attempts < 3
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId)any());
    heldContainer.incrementAssignmentAttempts();
    heldContainer.incrementAssignmentAttempts();
    // preemption - container released and resource asked again
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId)any());
    verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId4);
    verify(mockRMClient, times(4)).
    addContainerRequest(requestCaptor.capture());
    CookieContainerRequest reAdded = requestCaptor.getValue();
    Assert.assertEquals(pri6, reAdded.getPriority());
    Assert.assertEquals(taskAsk, reAdded.getCapability());
    Assert.assertEquals(mockTask3WaitCookie, reAdded.getCookie().getAppCookie());
    
    // remove fudging.
    scheduler.delayedContainerManager.delayedContainers.clear();
    
    scheduler.allocateTask(mockTask3Retry, taskAsk, null,
                           null, pri5, obj3, null);
    // no preemption - higher pri. exact match
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).releaseAssignedContainer((ContainerId)any());


    scheduler.allocateTask(mockTask2, taskAsk, null,
                           null, pri4, null, null);
    drainableAppCallback.drain();


    // mockTaskPri3Kill gets preempted
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).releaseAssignedContainer((ContainerId)any());
    verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId3);


    AppFinalStatus finalStatus =
        new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, "", appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.stop();
    drainableAppCallback.drain();
    scheduler.close();
  }

View Full Code Here

    TezAMRMClientAsync<CookieContainerRequest> amrmClient =
      mock(TezAMRMClientAsync.class);
    AppContext mockAppContext = mock(AppContext.class);
    when(mockAppContext.getAMState()).thenReturn(DAGAppMasterState.RUNNING);


    TaskSchedulerWithDrainableAppCallback taskScheduler =
      new TaskSchedulerWithDrainableAppCallback(
        appClient, new AlwaysMatchesContainerMatcher(), "host", 0, "",
        amrmClient, mockAppContext);
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler
        .getDrainableAppCallback();
    
    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
    taskScheduler.init(conf);
    
    RegisterApplicationMasterResponse mockRegResponse = mock(RegisterApplicationMasterResponse.class);
    Resource mockMaxResource = mock(Resource.class);
    Map<ApplicationAccessType, String> mockAcls = mock(Map.class);
    when(mockRegResponse.getMaximumResourceCapability()).thenReturn(
        mockMaxResource);
    when(mockRegResponse.getApplicationACLs()).thenReturn(mockAcls);
    when(amrmClient.registerApplicationMaster(anyString(), anyInt(), anyString()))
        .thenReturn(mockRegResponse);
    
    taskScheduler.start();
    
    Resource resource = Resource.newInstance(1024, 1);
    Priority priority = Priority.newInstance(1);


    String hostsTask1[] = { "host1" };
    String hostsTask2[] = { "non-allocated-host" };


    String defaultRack[] = { "/default-rack" };
    String otherRack[] = { "/other-rack" };


    Object mockTask1 = mock(Object.class);
    CookieContainerRequest mockCookie1 = mock(CookieContainerRequest.class,
        RETURNS_DEEP_STUBS);
    when(mockCookie1.getCookie().getTask()).thenReturn(mockTask1);


    Object mockTask2 = mock(Object.class);
    CookieContainerRequest mockCookie2 = mock(CookieContainerRequest.class,
        RETURNS_DEEP_STUBS);
    when(mockCookie2.getCookie().getTask()).thenReturn(mockTask2);


    Container containerHost1 = createContainer(1, "host1", resource, priority);
    Container containerHost3 = createContainer(2, "host3", resource, priority);
    List<Container> allocatedContainers = new LinkedList<Container>();


    allocatedContainers.add(containerHost3);
    allocatedContainers.add(containerHost1);


    final Map<String, List<CookieContainerRequest>> matchingMap = new HashMap<String, List<CookieContainerRequest>>();
    taskScheduler.allocateTask(mockTask1, resource, hostsTask1, defaultRack,
        priority, null, mockCookie1);
    drainableAppCallback.drain();


    List<CookieContainerRequest> host1List = new ArrayList<CookieContainerRequest>();
    host1List.add(mockCookie1);
    List<CookieContainerRequest> defaultRackList = new ArrayList<CookieContainerRequest>();
    defaultRackList.add(mockCookie1);
    matchingMap.put(hostsTask1[0], host1List);
    matchingMap.put(defaultRack[0], defaultRackList);


    List<CookieContainerRequest> nonAllocatedHostList = new ArrayList<TaskScheduler.CookieContainerRequest>();
    nonAllocatedHostList.add(mockCookie2);
    List<CookieContainerRequest> otherRackList = new ArrayList<TaskScheduler.CookieContainerRequest>();
    otherRackList.add(mockCookie2);
    taskScheduler.allocateTask(mockTask2, resource, hostsTask2, otherRack,
        priority, null, mockCookie2);
    drainableAppCallback.drain();
    matchingMap.put(hostsTask2[0], nonAllocatedHostList);
    matchingMap.put(otherRack[0], otherRackList);


    List<CookieContainerRequest> anyList = new LinkedList<TaskScheduler.CookieContainerRequest>();
    anyList.add(mockCookie1);
    anyList.add(mockCookie2);


    matchingMap.put(ResourceRequest.ANY, anyList);


    final List<ArrayList<CookieContainerRequest>> emptyList = new LinkedList<ArrayList<CookieContainerRequest>>();


    when(
        amrmClient.getMatchingRequests((Priority) any(), anyString(),
            (Resource) any())).thenAnswer(
        new Answer<List<? extends Collection<CookieContainerRequest>>>() {


          @Override
          public List<? extends Collection<CookieContainerRequest>> answer(
              InvocationOnMock invocation) throws Throwable {
            String location = (String) invocation.getArguments()[1];
            if (matchingMap.get(location) != null) {
              CookieContainerRequest matched = matchingMap.get(location).get(0);
              // Remove matched from matchingMap - assuming TaskScheduler will
              // pick the first entry.
              Iterator<Entry<String, List<CookieContainerRequest>>> iter = matchingMap
                  .entrySet().iterator();
              while (iter.hasNext()) {
                Entry<String, List<CookieContainerRequest>> entry = iter.next();
                if (entry.getValue().remove(matched)) {
                  if (entry.getValue().size() == 0) {
                    iter.remove();
                  }
                }
              }
              return Collections.singletonList(Collections
                  .singletonList(matched));
            } else {
              return emptyList;
            }
          }
        });


    taskScheduler.onContainersAllocated(allocatedContainers);
    drainableAppCallback.drain();


    ArgumentCaptor<Object> taskCaptor = ArgumentCaptor.forClass(Object.class);
    ArgumentCaptor<Container> containerCaptor = ArgumentCaptor
        .forClass(Container.class);
    verify(appClient, times(2)).taskAllocated(taskCaptor.capture(), any(),
        containerCaptor.capture());


    // Expected containerHost1 allocated to task1 due to locality,
    // containerHost3 allocated to task2.


    List<Container> assignedContainers = containerCaptor.getAllValues();
    int container1Pos = assignedContainers.indexOf(containerHost1);
    assertTrue("Container: " + containerHost1 + " was not assigned",
        container1Pos != -1);
    assertEquals("Task 1 was not allocated to containerHost1", mockTask1,
        taskCaptor.getAllValues().get(container1Pos));


    int container2Pos = assignedContainers.indexOf(containerHost3);
    assertTrue("Container: " + containerHost3 + " was not assigned",
        container2Pos != -1);
    assertEquals("Task 2 was not allocated to containerHost3", mockTask2,
        taskCaptor.getAllValues().get(container2Pos));


    AppFinalStatus finalStatus = new AppFinalStatus(
        FinalApplicationStatus.SUCCEEDED, "", "");
    when(appClient.getFinalAppStatus()).thenReturn(finalStatus);
    taskScheduler.close();
  }

View Full Code Here

0 1 2 3 4

TOP

Related Classes of org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableAppCallback

org.apache.tez.dag.app.rm.TestContainerReuse

org.apache.tez.dag.app.rm.TestTaskScheduler

All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.