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

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

Introduction

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

Prototype

@Public
    @Stable
    public static Resource newInstance(long memory, int vCores) 

Source Link

Usage

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

License:Apache License

/**
 * Normal flow of TaskAttempt//from   w w w . j a v a  2  s . c  o  m
 */
@Test(timeout = 5000)
public void testDeallocationBeforeAllocation() {
    MockLocalTaskSchedulerSerivce taskSchedulerService = new MockLocalTaskSchedulerSerivce(
            mock(TaskSchedulerAppCallback.class), mock(ContainerSignatureMatcher.class), "", 0, "",
            mock(AppContext.class));
    taskSchedulerService.init(new Configuration());
    taskSchedulerService.start();

    Task task = mock(Task.class);
    taskSchedulerService.allocateTask(task, Resource.newInstance(1024, 1), null, null, Priority.newInstance(1),
            null, null);
    taskSchedulerService.deallocateTask(task, false);
    // start the RequestHandler, DeallocateTaskRequest has higher priority, so will be processed first
    taskSchedulerService.startRequestHandlerThread();

    MockAsyncDelegateRequestHandler requestHandler = taskSchedulerService.getRequestHandler();
    requestHandler.drainRequest(1);
    assertEquals(1, requestHandler.deallocateCount);
    // The corresponding AllocateTaskRequest will be removed, so won't been processed.
    assertEquals(0, requestHandler.allocateCount);
    taskSchedulerService.stop();
}

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

License:Apache License

/**
 * TaskAttempt Killed from START_WAIT//w w  w.j av  a2 s .  c om
 */
@Test(timeout = 5000)
public void testDeallocationAfterAllocation() {
    MockLocalTaskSchedulerSerivce taskSchedulerService = new MockLocalTaskSchedulerSerivce(
            mock(TaskSchedulerAppCallback.class), mock(ContainerSignatureMatcher.class), "", 0, "",
            mock(AppContext.class));
    taskSchedulerService.init(new Configuration());
    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);
    requestHandler.drainRequest(2);
    assertEquals(1, requestHandler.deallocateCount);
    assertEquals(1, requestHandler.allocateCount);
    taskSchedulerService.stop();
}

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

License:Apache License

@SuppressWarnings({ "unchecked", "rawtypes" })
@Test(timeout = 5000)/*from  www . jav  a2 s . com*/
public void testTaskSchedulerPreemption() throws Exception {
    RackResolver.init(new YarnConfiguration());
    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);
    AppContext mockAppContext = mock(AppContext.class);
    when(mockAppContext.getAMState()).thenReturn(DAGAppMasterState.RUNNING);

    TezAMRMClientAsync<CookieContainerRequest> mockRMClient = 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 mockTask3KillA = mock(Object.class);
    Object mockTask3KillB = 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(mockTask3KillA, taskAsk, null, null, pri6, obj3, null);
    drainableAppCallback.drain();
    verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
    anyContainers.add(requestCaptor.getValue());
    // later one in the allocation gets killed between the two task3's
    scheduler.allocateTask(mockTask3KillB, taskAsk, null, null, pri6, obj3, null);
    drainableAppCallback.drain();
    verify(mockRMClient, times(4)).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 mockContainer3A = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer3A.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer3A.getResource()).thenReturn(taskAsk);
    when(mockContainer3A.getPriority()).thenReturn(pri6);
    ContainerId mockCId3A = mock(ContainerId.class);
    when(mockContainer3A.getId()).thenReturn(mockCId3A);
    containers.add(mockContainer3A);
    Container mockContainer3B = mock(Container.class, RETURNS_DEEP_STUBS);
    when(mockContainer3B.getNodeId().getHost()).thenReturn("host1");
    when(mockContainer3B.getResource()).thenReturn(taskAsk);
    when(mockContainer3B.getPriority()).thenReturn(pri2); // high priority container 
    ContainerId mockCId3B = mock(ContainerId.class);
    when(mockContainer3B.getId()).thenReturn(mockCId3B);
    containers.add(mockContainer3B);
    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(4, scheduler.taskAllocations.size());
    Assert.assertEquals(4096, scheduler.allocatedResources.getMemory());
    Assert.assertEquals(mockCId1, scheduler.taskAllocations.get(mockTask1).getId());
    Assert.assertEquals(mockCId2, scheduler.taskAllocations.get(mockTask3).getId());
    Assert.assertEquals(mockCId3A, scheduler.taskAllocations.get(mockTask3KillA).getId());
    // high priority container assigned to lower pri task. This task should still be preempted 
    // because the task priority is relevant for preemption and not the container priority
    Assert.assertEquals(mockCId3B, scheduler.taskAllocations.get(mockTask3KillB).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(5)).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());

    for (int i = 0; i < 11; ++i) {
        scheduler.allocateTask(mockTask2, taskAsk, null, null, pri4, null, null);
    }
    drainableAppCallback.drain();

    // mockTaskPri3KillB gets preempted to clear 10% of outstanding running preemptable tasks
    // this is also a higher priority container than the pending task priority but was running a 
    // lower priority task. Task priority is relevant for preemption and not container priority as
    // containers can run tasks of different priorities
    scheduler.getProgress();
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).releaseAssignedContainer((ContainerId) any());
    verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId3B);
    // next 3 heartbeats do nothing, waiting for the RM to act on the last released resources
    scheduler.getProgress();
    scheduler.getProgress();
    scheduler.getProgress();
    verify(mockRMClient, times(2)).releaseAssignedContainer((ContainerId) any());
    scheduler.getProgress();
    drainableAppCallback.drain();
    // Next oldest mockTaskPri3KillA gets preempted to clear 10% of outstanding running preemptable tasks
    verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());
    verify(mockRMClient, times(1)).releaseAssignedContainer(mockCId3A);

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

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

License:Apache License

@SuppressWarnings("unchecked")
@Test(timeout = 5000)// w w  w .  j a  v a  2 s  .c  om
public void testLocalityMatching() throws Exception {

    RackResolver.init(new Configuration());
    TaskSchedulerAppCallback appClient = mock(TaskSchedulerAppCallback.class);
    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<YarnTaskSchedulerService.CookieContainerRequest>();
    nonAllocatedHostList.add(mockCookie2);
    List<CookieContainerRequest> otherRackList = new ArrayList<YarnTaskSchedulerService.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<YarnTaskSchedulerService.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();
}

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

License:Apache License

@Test(timeout = 5000)
public void testTaskBasedAffinity() throws Exception {
    Configuration conf = new Configuration(false);
    schedulerHandler.init(conf);//from  w  w w.  ja v  a 2s.c  o  m
    schedulerHandler.start();

    TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class);
    TezTaskAttemptID taId = mock(TezTaskAttemptID.class);
    String affVertexName = "srcVertex";
    int affTaskIndex = 1;
    TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(affVertexName, affTaskIndex);
    VertexImpl affVertex = mock(VertexImpl.class);
    TaskImpl affTask = mock(TaskImpl.class);
    TaskAttemptImpl affAttempt = mock(TaskAttemptImpl.class);
    ContainerId affCId = mock(ContainerId.class);
    when(affVertex.getTotalTasks()).thenReturn(2);
    when(affVertex.getTask(affTaskIndex)).thenReturn(affTask);
    when(affTask.getSuccessfulAttempt()).thenReturn(affAttempt);
    when(affAttempt.getAssignedContainerID()).thenReturn(affCId);
    when(mockAppContext.getCurrentDAG().getVertex(affVertexName)).thenReturn(affVertex);
    Resource resource = Resource.newInstance(100, 1);
    AMSchedulerEventTALaunchRequest event = new AMSchedulerEventTALaunchRequest(taId, resource, null,
            mockTaskAttempt, locHint, 3, null);
    schedulerHandler.notify.set(false);
    schedulerHandler.handle(event);
    synchronized (schedulerHandler.notify) {
        while (!schedulerHandler.notify.get()) {
            schedulerHandler.notify.wait();
        }
    }

    // verify mockTaskAttempt affinitized to expected affCId
    verify(mockTaskScheduler, times(1)).allocateTask(mockTaskAttempt, resource, affCId, Priority.newInstance(3),
            null, event);

    schedulerHandler.stop();
    schedulerHandler.close();
}

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

License:Apache License

@Test(timeout = 5000)
public void testSimpleAllocate() throws Exception {
    Configuration conf = new Configuration(false);
    schedulerHandler.init(conf);//from w ww .j a  va  2  s . co  m
    schedulerHandler.start();

    TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class);
    TezTaskAttemptID mockAttemptId = mock(TezTaskAttemptID.class);
    when(mockAttemptId.getId()).thenReturn(0);
    when(mockTaskAttempt.getID()).thenReturn(mockAttemptId);
    Resource resource = Resource.newInstance(1024, 1);
    ContainerContext containerContext = new ContainerContext(new HashMap<String, LocalResource>(),
            new Credentials(), new HashMap<String, String>(), "");
    int priority = 10;
    TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(new HashSet<String>(), null);

    ContainerId mockCId = mock(ContainerId.class);
    Container container = mock(Container.class);
    when(container.getId()).thenReturn(mockCId);

    AMContainer mockAMContainer = mock(AMContainer.class);
    when(mockAMContainer.getContainerId()).thenReturn(mockCId);
    when(mockAMContainer.getState()).thenReturn(AMContainerState.IDLE);

    when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer);

    AMSchedulerEventTALaunchRequest lr = new AMSchedulerEventTALaunchRequest(mockAttemptId, resource, null,
            mockTaskAttempt, locHint, priority, containerContext, 0, 0, 0);
    schedulerHandler.taskAllocated(0, mockTaskAttempt, lr, container);
    assertEquals(1, mockEventHandler.events.size());
    assertTrue(mockEventHandler.events.get(0) instanceof AMContainerEventAssignTA);
    AMContainerEventAssignTA assignEvent = (AMContainerEventAssignTA) mockEventHandler.events.get(0);
    assertEquals(priority, assignEvent.getPriority());
    assertEquals(mockAttemptId, assignEvent.getTaskAttemptId());
}

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

License:Apache License

@Test(timeout = 5000)
public void testTaskBasedAffinity() throws Exception {
    Configuration conf = new Configuration(false);
    schedulerHandler.init(conf);//  ww  w. j  av  a 2s.  c o  m
    schedulerHandler.start();

    TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class);
    TezTaskAttemptID taId = mock(TezTaskAttemptID.class);
    String affVertexName = "srcVertex";
    int affTaskIndex = 1;
    TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(affVertexName, affTaskIndex);
    VertexImpl affVertex = mock(VertexImpl.class);
    TaskImpl affTask = mock(TaskImpl.class);
    TaskAttemptImpl affAttempt = mock(TaskAttemptImpl.class);
    ContainerId affCId = mock(ContainerId.class);
    when(affVertex.getTotalTasks()).thenReturn(2);
    when(affVertex.getTask(affTaskIndex)).thenReturn(affTask);
    when(affTask.getSuccessfulAttempt()).thenReturn(affAttempt);
    when(affAttempt.getAssignedContainerID()).thenReturn(affCId);
    when(mockAppContext.getCurrentDAG().getVertex(affVertexName)).thenReturn(affVertex);
    Resource resource = Resource.newInstance(100, 1);
    AMSchedulerEventTALaunchRequest event = new AMSchedulerEventTALaunchRequest(taId, resource, null,
            mockTaskAttempt, locHint, 3, null, 0, 0, 0);
    schedulerHandler.notify.set(false);
    schedulerHandler.handle(event);
    synchronized (schedulerHandler.notify) {
        while (!schedulerHandler.notify.get()) {
            schedulerHandler.notify.wait();
        }
    }

    // verify mockTaskAttempt affinitized to expected affCId
    verify(mockTaskScheduler, times(1)).allocateTask(mockTaskAttempt, resource, affCId, Priority.newInstance(3),
            null, event);

    schedulerHandler.stop();
    schedulerHandler.close();
}

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

License:Apache License

@Test(timeout = 5000)
public void testTaskSchedulerRouting() throws Exception {
    Configuration conf = new Configuration(false);
    UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf);

    String customSchedulerName = "fakeScheduler";
    List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>();
    ByteBuffer bb = ByteBuffer.allocate(4);
    bb.putInt(0, 3);//from   w  ww  .ja v a 2  s. c o m
    UserPayload userPayload = UserPayload.create(bb);
    taskSchedulers.add(new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName())
            .setUserPayload(userPayload));
    taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null)
            .setUserPayload(defaultPayload));

    TSEHForMultipleSchedulersTest tseh = new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService,
            mockEventHandler, mockSigMatcher, mockWebUIService, taskSchedulers, false);

    tseh.init(conf);
    tseh.start();

    // Verify that the YARN task scheduler is installed by default
    assertTrue(tseh.getYarnSchedulerCreated());
    assertFalse(tseh.getUberSchedulerCreated());
    assertEquals(2, tseh.getNumCreateInvocations());

    // Verify the order of the schedulers
    assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0));
    assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1));

    verify(tseh.getTestTaskScheduler(0)).initialize();
    verify(tseh.getTestTaskScheduler(0)).start();

    ApplicationId appId = ApplicationId.newInstance(1000, 1);
    TezDAGID dagId = TezDAGID.getInstance(appId, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagId, 1);
    TezTaskID taskId1 = TezTaskID.getInstance(vertexID, 1);
    TezTaskAttemptID attemptId11 = TezTaskAttemptID.getInstance(taskId1, 1);
    TezTaskID taskId2 = TezTaskID.getInstance(vertexID, 2);
    TezTaskAttemptID attemptId21 = TezTaskAttemptID.getInstance(taskId2, 1);

    Resource resource = Resource.newInstance(1024, 1);

    TaskAttempt mockTaskAttempt1 = mock(TaskAttempt.class);
    TaskAttempt mockTaskAttempt2 = mock(TaskAttempt.class);

    AMSchedulerEventTALaunchRequest launchRequest1 = new AMSchedulerEventTALaunchRequest(attemptId11, resource,
            mock(TaskSpec.class), mockTaskAttempt1, mock(TaskLocationHint.class), 1,
            mock(ContainerContext.class), 0, 0, 0);

    tseh.handle(launchRequest1);

    verify(tseh.getTestTaskScheduler(0)).allocateTask(eq(mockTaskAttempt1), eq(resource), any(String[].class),
            any(String[].class), any(Priority.class), any(Object.class), eq(launchRequest1));

    AMSchedulerEventTALaunchRequest launchRequest2 = new AMSchedulerEventTALaunchRequest(attemptId21, resource,
            mock(TaskSpec.class), mockTaskAttempt2, mock(TaskLocationHint.class), 1,
            mock(ContainerContext.class), 1, 0, 0);
    tseh.handle(launchRequest2);
    verify(tseh.getTestTaskScheduler(1)).allocateTask(eq(mockTaskAttempt2), eq(resource), any(String[].class),
            any(String[].class), any(Priority.class), any(Object.class), eq(launchRequest2));
}

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

License:Apache License

@Test(timeout = 10000)
public void testMatchingRequestsForTopPriority() {
    String[] hosts = { "host1" };
    String[] racks = { "rack1" };
    AMRMClient.ContainerRequest req1 = new AMRMClient.ContainerRequest(Resource.newInstance(2048, 1), hosts,
            racks, Priority.newInstance(1));
    AMRMClient.ContainerRequest req2 = new AMRMClient.ContainerRequest(Resource.newInstance(1024, 1), hosts,
            racks, Priority.newInstance(2));
    AMRMClient.ContainerRequest req3 = new AMRMClient.ContainerRequest(Resource.newInstance(1024, 1), hosts,
            racks, Priority.newInstance(3));
    amrmClient.addContainerRequest(req1);
    amrmClient.addContainerRequest(req2);
    amrmClient.addContainerRequest(req3);

    Assert.assertTrue(// www .ja  v a 2  s . co m
            amrmClient.getMatchingRequestsForTopPriority("host1", Resource.newInstance(1024, 1)).isEmpty());

    List<? extends Collection<AMRMClient.ContainerRequest>> ret = amrmClient
            .getMatchingRequestsForTopPriority("host1", Resource.newInstance(2048, 1));
    Assert.assertFalse(ret.isEmpty());
    Assert.assertEquals(req1, ret.get(0).iterator().next());

    amrmClient.removeContainerRequest(req1);

    ret = amrmClient.getMatchingRequestsForTopPriority("host1", Resource.newInstance(1024, 1));
    Assert.assertFalse(ret.isEmpty());
    Assert.assertEquals(req2, ret.get(0).iterator().next());
}

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

License:Apache License

public TezTestServiceTaskSchedulerService(TaskSchedulerContext taskSchedulerContext) {
    // Accepting configuration here to allow setting up fields as final
    super(taskSchedulerContext);
    this.serviceHosts = new LinkedList<String>();
    this.containerFactory = new ContainerFactory(taskSchedulerContext.getApplicationAttemptId(),
            taskSchedulerContext.getCustomClusterIdentifier());

    Configuration conf = null;//from  w  w w .j  a  va 2  s  .co  m
    try {
        conf = TezUtils.createConfFromUserPayload(taskSchedulerContext.getInitialUserPayload());
    } catch (IOException e) {
        throw new TezUncheckedException(e);
    }
    this.memoryPerInstance = conf.getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_MEMORY_PER_INSTANCE_MB,
            -1);
    Preconditions.checkArgument(memoryPerInstance > 0,
            TezTestServiceConfConstants.TEZ_TEST_SERVICE_MEMORY_PER_INSTANCE_MB + " must be configured");

    this.executorsPerInstance = conf
            .getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_NUM_EXECUTORS_PER_INSTANCE, -1);
    Preconditions.checkArgument(executorsPerInstance > 0,
            TezTestServiceConfConstants.TEZ_TEST_SERVICE_NUM_EXECUTORS_PER_INSTANCE + " must be configured");

    this.coresPerInstance = conf.getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_VCPUS_PER_INSTANCE,
            executorsPerInstance);

    this.containerPort = conf.getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_RPC_PORT, -1);
    Preconditions.checkArgument(executorsPerInstance > 0,
            TezTestServiceConfConstants.TEZ_TEST_SERVICE_RPC_PORT + " must be configured");

    int memoryPerContainer = (int) (memoryPerInstance / (float) executorsPerInstance);
    int coresPerContainer = (int) (coresPerInstance / (float) executorsPerInstance);
    this.resourcePerContainer = Resource.newInstance(memoryPerContainer, coresPerContainer);

    String[] hosts = conf.getTrimmedStrings(TezTestServiceConfConstants.TEZ_TEST_SERVICE_HOSTS);
    if (hosts == null || hosts.length == 0) {
        hosts = new String[] { "localhost" };
    }
    for (String host : hosts) {
        serviceHosts.add(host);
    }

    LOG.info("Running with configuration: " + "memoryPerInstance=" + memoryPerInstance + ", vcoresPerInstance="
            + coresPerInstance + ", executorsPerInstance=" + executorsPerInstance
            + ", resourcePerContainerInferred=" + resourcePerContainer + ", hosts=" + serviceHosts.toString());

}