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

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

Introduction

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

Prototype

@Public
    @Stable
    public static Priority newInstance(int p) 

Source Link

Usage

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

License:Apache License

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

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

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NEW_CONTAINERS_ENABLED, true);
    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();/* ww w.j a  va 2  s . com*/
    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 node1Rack1 = "n1r1";
    final String node2Rack1 = "n2r1";
    final String node1Rack2 = "n1r2";
    MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);

    Priority priorityv0 = Priority.newInstance(1);
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);

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

    // verify idle container is kept for now
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    clock.incrementTime(2000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(cid1);
    verify(mockRMClient, never()).releaseAssignedContainer(cid1);

    // verify idle container is released without being assigned to a task because rack-local reuse is
    // disabled
    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient, never()).removeContainerRequest(reqv0t0);
    verify(mockApp, never()).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    // cancel the task request
    assertFalse(scheduler.deallocateTask(taskv0t0.task, false, null, null));

    // allocate another container that's node-local
    ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
    NodeId n1r1 = NodeId.newInstance(node1Rack1, 1);
    Container container2 = Container.newInstance(cid2, n1r1, null, taskv0t0.capability, priorityv0, null);
    scheduler.onContainersAllocated(Collections.singletonList(container2));
    clock.incrementTime(2000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(cid2);
    verify(mockRMClient, never()).releaseAssignedContainer(cid2);

    // reschedule the task, verify it's now scheduled without a container request
    // since node-local idle container is available
    reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0, false);
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container2);
    verify(mockRMClient).removeContainerRequest(reqv0t0);
}

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

License:Apache License

@Test(timeout = 5000)
public void testSimpleLocalAllocation() throws IOException, InterruptedException {

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper();

    try {/*from   w  w w  . j a v  a2 s .  c o m*/
        Priority priority1 = Priority.newInstance(1);
        String[] hosts1 = new String[] { HOST1 };

        Object task1 = new Object();
        Object clientCookie1 = new Object();

        tsWrapper.controlScheduler(true);
        tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1);

        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();

        verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class));
        // TODO Verify this is on host1.
        assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

@Test(timeout = 5000)
public void testSimpleNoLocalityAllocation() throws IOException, InterruptedException {
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper();

    try {/*from  www .  j  av  a  2 s  . c o m*/
        Priority priority1 = Priority.newInstance(1);

        Object task1 = new Object();
        Object clientCookie1 = new Object();
        tsWrapper.controlScheduler(true);
        tsWrapper.allocateTask(task1, null, priority1, clientCookie1);
        tsWrapper.signalSchedulerRun();
        tsWrapper.awaitSchedulerRun();
        verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class));
        assertEquals(1, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

@Test(timeout = 5000)
public void testPreemption() throws InterruptedException, IOException {

    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);
    String[] hosts = new String[] { HOST1 };
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1);
    try {/*from ww  w. j a  va  2 s .c  o m*/

        Object task1 = "task1";
        Object clientCookie1 = "cookie1";
        Object task2 = "task2";
        Object clientCookie2 = "cookie2";
        Object task3 = "task3";
        Object clientCookie3 = "cookie3";
        Object task4 = "task4";
        Object clientCookie4 = "cookie4";

        tsWrapper.controlScheduler(true);
        tsWrapper.allocateTask(task1, hosts, priority2, clientCookie1);
        tsWrapper.allocateTask(task2, hosts, priority2, clientCookie2);
        tsWrapper.allocateTask(task3, hosts, priority2, clientCookie3);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numLocalAllocations == 2) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(any(Object.class), any(Object.class),
                any(Container.class));
        assertEquals(2, tsWrapper.ts.dagStats.numLocalAllocations);
        assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);

        reset(tsWrapper.mockAppCallback);

        tsWrapper.allocateTask(task4, hosts, priority1, clientCookie4);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class));

        tsWrapper.deallocateTask(task2, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4), eq(clientCookie4),
                any(Container.class));

    } finally {
        tsWrapper.shutdown();
    }

}

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

License:Apache License

@Test(timeout = 5000)
public void testNodeDisabled() throws IOException, InterruptedException {
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(10000l);
    try {//from   w w w  . j  av a 2 s .  com
        Priority priority1 = Priority.newInstance(1);
        String[] hosts1 = new String[] { HOST1 };
        Object task1 = new Object();
        Object clientCookie1 = new Object();
        tsWrapper.controlScheduler(true);
        tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 1) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class));
        assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
        assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
        assertEquals(0, tsWrapper.ts.dagStats.numNonLocalAllocations);
        assertEquals(1, tsWrapper.ts.dagStats.numTotalAllocations);

        tsWrapper.resetAppCallback();

        tsWrapper.clock.setTime(10000l);
        tsWrapper.rejectExecution(task1);

        // Verify that the node is blacklisted
        assertEquals(1, tsWrapper.ts.dagStats.numRejectedTasks);
        assertEquals(3, tsWrapper.ts.instanceToNodeMap.size());
        LlapTaskSchedulerService.NodeInfo disabledNodeInfo = tsWrapper.ts.disabledNodesQueue.peek();
        assertNotNull(disabledNodeInfo);
        assertEquals(HOST1, disabledNodeInfo.serviceInstance.getHost());
        assertEquals((10000l), disabledNodeInfo.getDelay(TimeUnit.MILLISECONDS));
        assertEquals((10000l + 10000l), disabledNodeInfo.expireTimeMillis);

        Object task2 = new Object();
        Object clientCookie2 = new Object();
        tsWrapper.allocateTask(task2, hosts1, priority1, clientCookie2);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 2) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback).taskAllocated(eq(task2), eq(clientCookie2), any(Container.class));
        assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations);
        assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
        assertEquals(1, tsWrapper.ts.dagStats.numNonLocalAllocations);
        assertEquals(2, tsWrapper.ts.dagStats.numTotalAllocations);

    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

@Test(timeout = 5000)
public void testNodeReEnabled() throws InterruptedException, IOException {
    // Based on actual timing.
    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(1000l);
    try {/*  w  w  w . ja v a 2  s. c o m*/
        Priority priority1 = Priority.newInstance(1);
        String[] hosts1 = new String[] { HOST1 };
        String[] hosts2 = new String[] { HOST2 };
        String[] hosts3 = new String[] { HOST3 };

        Object task1 = new Object();
        Object clientCookie1 = new Object();
        Object task2 = new Object();
        Object clientCookie2 = new Object();
        Object task3 = new Object();
        Object clientCookie3 = new Object();

        tsWrapper.controlScheduler(true);
        tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1);
        tsWrapper.allocateTask(task2, hosts2, priority1, clientCookie2);
        tsWrapper.allocateTask(task3, hosts3, priority1, clientCookie3);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback, times(3)).taskAllocated(any(Object.class), any(Object.class),
                any(Container.class));
        assertEquals(3, tsWrapper.ts.dagStats.numLocalAllocations);
        assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
        assertEquals(3, tsWrapper.ts.dagStats.numTotalAllocations);

        tsWrapper.resetAppCallback();

        tsWrapper.rejectExecution(task1);
        tsWrapper.rejectExecution(task2);
        tsWrapper.rejectExecution(task3);

        // Verify that the node is blacklisted
        assertEquals(3, tsWrapper.ts.dagStats.numRejectedTasks);
        assertEquals(3, tsWrapper.ts.instanceToNodeMap.size());
        assertEquals(3, tsWrapper.ts.disabledNodesQueue.size());

        Object task4 = new Object();
        Object clientCookie4 = new Object();
        Object task5 = new Object();
        Object clientCookie5 = new Object();
        Object task6 = new Object();
        Object clientCookie6 = new Object();
        tsWrapper.allocateTask(task4, hosts1, priority1, clientCookie4);
        tsWrapper.allocateTask(task5, hosts2, priority1, clientCookie5);
        tsWrapper.allocateTask(task6, hosts3, priority1, clientCookie6);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 6) {
                break;
            }
        }

        ArgumentCaptor<Container> argumentCaptor = ArgumentCaptor.forClass(Container.class);
        verify(tsWrapper.mockAppCallback, times(3)).taskAllocated(any(Object.class), any(Object.class),
                argumentCaptor.capture());

        // which affects the locality matching
        assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest);
        assertEquals(6, tsWrapper.ts.dagStats.numTotalAllocations);

    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

private void forceLocalityTest1(boolean forceLocality) throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hosts = new String[] { HOST1, HOST2 };

    String[] hostsH1 = new String[] { HOST1 };
    String[] hostsH2 = new String[] { HOST2 };

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1,
            (forceLocality ? -1l : 0l));

    try {//from   w w w .  j a  va 2 s. c  om
        Object task1 = "task1";
        Object clientCookie1 = "cookie1";
        Object task2 = "task2";
        Object clientCookie2 = "cookie2";
        Object task3 = "task3";
        Object clientCookie3 = "cookie3";
        Object task4 = "task4";
        Object clientCookie4 = "cookie4";
        Object task5 = "task5";
        Object clientCookie5 = "cookie5";

        tsWrapper.controlScheduler(true);
        //H1 - should allocate
        tsWrapper.allocateTask(task1, hostsH1, priority1, clientCookie1);
        //H1 - should allocate
        tsWrapper.allocateTask(task2, hostsH1, priority1, clientCookie2);
        //H1 - no capacity if force, should allocate otherwise
        tsWrapper.allocateTask(task3, hostsH1, priority1, clientCookie3);
        //H2 - should allocate
        tsWrapper.allocateTask(task4, hostsH2, priority1, clientCookie4);
        //No location - should allocate if force, no capacity otherwise
        tsWrapper.allocateTask(task5, null, priority1, clientCookie5);

        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 4) {
                break;
            }
        }

        // Verify no preemption requests - since everything is at the same priority
        verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
        ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
        verify(tsWrapper.mockAppCallback, times(4)).taskAllocated(argumentCaptor.capture(), any(Object.class),
                any(Container.class));
        assertEquals(4, argumentCaptor.getAllValues().size());
        assertEquals(task1, argumentCaptor.getAllValues().get(0));
        assertEquals(task2, argumentCaptor.getAllValues().get(1));
        if (forceLocality) {
            // task3 not allocated
            assertEquals(task4, argumentCaptor.getAllValues().get(2));
            assertEquals(task5, argumentCaptor.getAllValues().get(3));
        } else {
            assertEquals(task3, argumentCaptor.getAllValues().get(2));
            assertEquals(task4, argumentCaptor.getAllValues().get(3));
        }

        //Complete one task on host1.
        tsWrapper.deallocateTask(task1, true, null);

        reset(tsWrapper.mockAppCallback);

        // Try scheduling again.
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 5) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
        argumentCaptor = ArgumentCaptor.forClass(Object.class);
        verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(argumentCaptor.capture(), any(Object.class),
                any(Container.class));
        assertEquals(1, argumentCaptor.getAllValues().size());
        if (forceLocality) {
            assertEquals(task3, argumentCaptor.getAllValues().get(0));
        } else {
            assertEquals(task5, argumentCaptor.getAllValues().get(0));
        }

    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

@Test(timeout = 5000)
public void testForcedLocalityUnknownHost() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);

    String[] hostsKnown = new String[] { HOST1 };
    String[] hostsUnknown = new String[] { HOST2 };

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hostsKnown, 1, 1,
            -1l);/*from   www .j a v  a  2s  .c  om*/
    try {
        Object task1 = "task1";
        Object clientCookie1 = "cookie1";

        Object task2 = "task2";
        Object clientCookie2 = "cookie2";

        tsWrapper.controlScheduler(true);
        // Should allocate since H2 is not known.
        tsWrapper.allocateTask(task1, hostsUnknown, priority1, clientCookie1);
        tsWrapper.allocateTask(task2, hostsKnown, priority1, clientCookie2);

        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 2) {
                break;
            }
        }

        ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
        verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(argumentCaptor.capture(), any(Object.class),
                any(Container.class));
        assertEquals(2, argumentCaptor.getAllValues().size());
        assertEquals(task1, argumentCaptor.getAllValues().get(0));
        assertEquals(task2, argumentCaptor.getAllValues().get(1));

    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

@Test(timeout = 5000)
public void testForcedLocalityPreemption() throws IOException, InterruptedException {
    Priority priority1 = Priority.newInstance(1);
    Priority priority2 = Priority.newInstance(2);
    String[] hosts = new String[] { HOST1, HOST2 };

    String[] hostsH1 = new String[] { HOST1 };
    String[] hostsH2 = new String[] { HOST2 };

    TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, -1l);

    // Fill up host1 with p2 tasks.
    // Leave host2 empty
    // Try running p1 task on host1 - should preempt

    try {/*from ww  w  .j av  a  2  s . c om*/
        Object task1 = "task1";
        Object clientCookie1 = "cookie1";
        Object task2 = "task2";
        Object clientCookie2 = "cookie2";
        Object task3 = "task3";
        Object clientCookie3 = "cookie3";
        Object task4 = "task4";
        Object clientCookie4 = "cookie4";

        tsWrapper.controlScheduler(true);
        tsWrapper.allocateTask(task1, hostsH1, priority2, clientCookie1);
        tsWrapper.allocateTask(task2, hostsH1, priority2, clientCookie2);
        // This request at a lower priority should not affect anything.
        tsWrapper.allocateTask(task3, hostsH1, priority2, clientCookie3);
        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numLocalAllocations == 2) {
                break;
            }
        }

        verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class));
        ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class);
        verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(argumentCaptor.capture(), any(Object.class),
                any(Container.class));
        assertEquals(2, argumentCaptor.getAllValues().size());
        assertEquals(task1, argumentCaptor.getAllValues().get(0));
        assertEquals(task2, argumentCaptor.getAllValues().get(1));

        reset(tsWrapper.mockAppCallback);
        // Allocate t4 at higher priority. t3 should not be allocated,
        // and a preemption should be attempted on host1, despite host2 having available capacity
        tsWrapper.allocateTask(task4, hostsH1, priority1, clientCookie4);

        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class));

        tsWrapper.deallocateTask(task1, false, TaskAttemptEndReason.INTERNAL_PREEMPTION);

        while (true) {
            tsWrapper.signalSchedulerRun();
            tsWrapper.awaitSchedulerRun();
            if (tsWrapper.ts.dagStats.numTotalAllocations == 3) {
                break;
            }
        }
        verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4), eq(clientCookie4),
                any(Container.class));

    } finally {
        tsWrapper.shutdown();
    }
}

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

License:Apache License

@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);

    LocalContainerFactory containerFactory = new LocalContainerFactory(createMockAppContext());
    HashMap<Object, Container> taskAllocations = new LinkedHashMap<Object, Container>();
    PriorityBlockingQueue<TaskRequest> taskRequestQueue = new PriorityBlockingQueue<TaskRequest>();
    TaskSchedulerAppCallback appClientDelegate = mock(TaskSchedulerAppCallback.class);

    // Object under test
    AsyncDelegateRequestHandler requestHandler = new AsyncDelegateRequestHandler(taskRequestQueue,
            containerFactory, taskAllocations, appClientDelegate, 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.processRequest();
    }//from w ww.  jav  a 2  s  .com

    // 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.shouldWait());

    // Deallocate down to zero
    for (int i = 0; i < MAX_TASKS; i++) {
        requestHandler.addDeallocateTaskRequest(new Long(i));
        requestHandler.processRequest();
    }

    // All allocated tasks should have been removed
    Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size());
}