Example usage for org.apache.hadoop.conf Configuration setLong

List of usage examples for org.apache.hadoop.conf Configuration setLong

Introduction

In this page you can find the example usage for org.apache.hadoop.conf Configuration setLong.

Prototype

public void setLong(String name, long value) 

Source Link

Document

Set the value of the name property to a long.

Usage

From source file:org.apache.tez.auxservices.TestShuffleHandlerJobs.java

License:Apache License

@BeforeClass
public static void setup() throws IOException {
    try {/*from  w ww .  j  a va  2s.  c  o  m*/
        conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 1);
        conf.setInt(YarnConfiguration.NM_CONTAINER_MGR_THREAD_COUNT, 22);
        dfsCluster = new MiniDFSCluster.Builder(conf).numDataNodes(NUM_DNS).format(true).build();
        remoteFs = dfsCluster.getFileSystem();
    } catch (IOException io) {
        throw new RuntimeException("problem starting mini dfs cluster", io);
    }

    if (!(new File(MiniTezCluster.APPJAR)).exists()) {
        LOG.info("MRAppJar " + MiniTezCluster.APPJAR + " not found. Not running test.");
        return;
    }

    if (tezCluster == null) {
        tezCluster = new MiniTezCluster(TestShuffleHandlerJobs.class.getName(), NUM_NMS, 1, 1);
        Configuration conf = new Configuration();
        conf.set(YarnConfiguration.NM_AUX_SERVICES, ShuffleHandler.TEZ_SHUFFLE_SERVICEID);
        String serviceStr = String.format(YarnConfiguration.NM_AUX_SERVICE_FMT,
                ShuffleHandler.TEZ_SHUFFLE_SERVICEID);
        conf.set(serviceStr, ShuffleHandler.class.getName());
        conf.setInt(ShuffleHandler.SHUFFLE_PORT_CONFIG_KEY, 0);
        conf.set("fs.defaultFS", remoteFs.getUri().toString()); // use HDFS
        conf.setLong(YarnConfiguration.DEBUG_NM_DELETE_DELAY_SEC, 0l);
        tezCluster.init(conf);
        tezCluster.start();
    }

}

From source file:org.apache.tez.common.TestTezUtils.java

License:Apache License

private Configuration getConf() {
    Configuration conf = new Configuration(false);
    conf.set("test1", "value1");
    conf.setBoolean("test2", true);
    conf.setDouble("test3", 1.2345);
    conf.setInt("test4", 34567);
    conf.setLong("test5", 1234567890L);
    conf.setStrings("test6", "S1", "S2", "S3");
    return conf;/*from   ww w. j a v  a2s.c o m*/
}

From source file:org.apache.tez.dag.app.dag.impl.TestVertexScheduler.java

License:Apache License

@SuppressWarnings({ "unchecked", "rawtypes" })
@Test(timeout = 5000)/*from   ww w .  java  2  s  .co  m*/
public void testShuffleVertexManagerAutoParallelism() throws IOException {
    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_SHUFFLE_VERTEX_MANAGER_ENABLE_AUTO_PARALLEL, true);
    conf.setLong(TezConfiguration.TEZ_AM_SHUFFLE_VERTEX_MANAGER_DESIRED_TASK_INPUT_SIZE, 1000L);
    ShuffleVertexManager scheduler = null;
    EventHandler mockEventHandler = mock(EventHandler.class);
    TezDAGID dagId = TezDAGID.getInstance("1", 1, 1);
    HashMap<Vertex, Edge> mockInputVertices = new HashMap<Vertex, Edge>();
    Vertex mockSrcVertex1 = mock(Vertex.class);
    TezVertexID mockSrcVertexId1 = TezVertexID.getInstance(dagId, 1);
    EdgeProperty eProp1 = new EdgeProperty(EdgeProperty.DataMovementType.SCATTER_GATHER,
            EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, new OutputDescriptor("out"),
            new InputDescriptor("in"));
    when(mockSrcVertex1.getVertexId()).thenReturn(mockSrcVertexId1);
    Vertex mockSrcVertex2 = mock(Vertex.class);
    TezVertexID mockSrcVertexId2 = TezVertexID.getInstance(dagId, 2);
    EdgeProperty eProp2 = new EdgeProperty(EdgeProperty.DataMovementType.SCATTER_GATHER,
            EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, new OutputDescriptor("out"),
            new InputDescriptor("in"));
    when(mockSrcVertex2.getVertexId()).thenReturn(mockSrcVertexId2);
    Vertex mockSrcVertex3 = mock(Vertex.class);
    TezVertexID mockSrcVertexId3 = TezVertexID.getInstance(dagId, 3);
    EdgeProperty eProp3 = new EdgeProperty(EdgeProperty.DataMovementType.BROADCAST,
            EdgeProperty.DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, new OutputDescriptor("out"),
            new InputDescriptor("in"));
    when(mockSrcVertex3.getVertexId()).thenReturn(mockSrcVertexId3);

    Vertex mockManagedVertex = mock(Vertex.class);
    TezVertexID mockManagedVertexId = TezVertexID.getInstance(dagId, 4);
    when(mockManagedVertex.getVertexId()).thenReturn(mockManagedVertexId);
    when(mockManagedVertex.getInputVertices()).thenReturn(mockInputVertices);

    mockInputVertices.put(mockSrcVertex1, new Edge(eProp1, mockEventHandler));
    mockInputVertices.put(mockSrcVertex2, new Edge(eProp2, mockEventHandler));
    mockInputVertices.put(mockSrcVertex3, new Edge(eProp3, mockEventHandler));

    // check initialization
    scheduler = createScheduler(conf, mockManagedVertex, 0.1f, 0.1f);
    Assert.assertTrue(scheduler.bipartiteSources.size() == 2);
    Assert.assertTrue(scheduler.bipartiteSources.containsKey(mockSrcVertexId1));
    Assert.assertTrue(scheduler.bipartiteSources.containsKey(mockSrcVertexId2));

    final HashMap<TezTaskID, Task> managedTasks = new HashMap<TezTaskID, Task>();
    final TezTaskID mockTaskId1 = TezTaskID.getInstance(mockManagedVertexId, 0);
    managedTasks.put(mockTaskId1, null);
    final TezTaskID mockTaskId2 = TezTaskID.getInstance(mockManagedVertexId, 1);
    managedTasks.put(mockTaskId2, null);
    final TezTaskID mockTaskId3 = TezTaskID.getInstance(mockManagedVertexId, 2);
    managedTasks.put(mockTaskId3, null);
    final TezTaskID mockTaskId4 = TezTaskID.getInstance(mockManagedVertexId, 3);
    managedTasks.put(mockTaskId4, null);

    when(mockManagedVertex.getTotalTasks()).thenReturn(managedTasks.size());
    when(mockManagedVertex.getTasks()).thenReturn(managedTasks);

    final HashSet<TezTaskID> scheduledTasks = new HashSet<TezTaskID>();
    doAnswer(new Answer() {
        public Object answer(InvocationOnMock invocation) {
            Object[] args = invocation.getArguments();
            scheduledTasks.clear();
            scheduledTasks.addAll((Collection<TezTaskID>) args[0]);
            return null;
        }
    }).when(mockManagedVertex).scheduleTasks(anyCollection());

    final Map<Vertex, EdgeManager> newEdgeManagers = new HashMap<Vertex, EdgeManager>();

    doAnswer(new Answer() {
        public Object answer(InvocationOnMock invocation) {
            managedTasks.remove(mockTaskId3);
            managedTasks.remove(mockTaskId4);
            newEdgeManagers.clear();
            newEdgeManagers.putAll((Map<Vertex, EdgeManager>) invocation.getArguments()[1]);
            return null;
        }
    }).when(mockManagedVertex).setParallelism(eq(2), anyMap());

    // source vertices have 0 tasks. immediate start of all managed tasks
    when(mockSrcVertex1.getTotalTasks()).thenReturn(0);
    when(mockSrcVertex2.getTotalTasks()).thenReturn(0);
    scheduler.onVertexStarted(null);
    Assert.assertTrue(scheduler.pendingTasks.isEmpty());
    Assert.assertTrue(scheduledTasks.size() == 4); // all tasks scheduled
    scheduledTasks.clear();

    when(mockSrcVertex1.getTotalTasks()).thenReturn(2);
    when(mockSrcVertex2.getTotalTasks()).thenReturn(2);

    TezTaskAttemptID mockSrcAttemptId11 = TezTaskAttemptID
            .getInstance(TezTaskID.getInstance(mockSrcVertexId1, 0), 0);
    TezTaskAttemptID mockSrcAttemptId12 = TezTaskAttemptID
            .getInstance(TezTaskID.getInstance(mockSrcVertexId1, 1), 0);
    TezTaskAttemptID mockSrcAttemptId21 = TezTaskAttemptID
            .getInstance(TezTaskID.getInstance(mockSrcVertexId2, 0), 0);
    TezTaskAttemptID mockSrcAttemptId31 = TezTaskAttemptID
            .getInstance(TezTaskID.getInstance(mockSrcVertexId3, 0), 0);

    byte[] payload = VertexManagerEventPayloadProto.newBuilder().setOutputSize(5000L).build().toByteArray();
    VertexManagerEvent vmEvent = new VertexManagerEvent("Vertex", payload);
    // parallelism not change due to large data size
    scheduler = createScheduler(conf, mockManagedVertex, 0.1f, 0.1f);
    scheduler.onVertexStarted(null);
    Assert.assertTrue(scheduler.pendingTasks.size() == 4); // no tasks scheduled
    Assert.assertTrue(scheduler.numSourceTasks == 4);
    scheduler.onVertexManagerEventReceived(vmEvent);
    scheduler.onSourceTaskCompleted(mockSrcAttemptId11);
    // managedVertex tasks reduced
    verify(mockManagedVertex, times(0)).setParallelism(anyInt(), anyMap());
    Assert.assertEquals(0, scheduler.pendingTasks.size()); // all tasks scheduled
    Assert.assertEquals(4, scheduledTasks.size());
    Assert.assertEquals(1, scheduler.numSourceTasksCompleted); // TODO
    Assert.assertEquals(5000L, scheduler.completedSourceTasksOutputSize);

    // parallelism changed due to small data size
    scheduledTasks.clear();
    payload = VertexManagerEventPayloadProto.newBuilder().setOutputSize(500L).build().toByteArray();
    vmEvent = new VertexManagerEvent("Vertex", payload);

    scheduler = createScheduler(conf, mockManagedVertex, 0.5f, 0.5f);
    scheduler.onVertexStarted(null);
    Assert.assertEquals(4, scheduler.pendingTasks.size()); // no tasks scheduled
    Assert.assertEquals(4, scheduler.numSourceTasks);
    // task completion from non-bipartite stage does nothing
    scheduler.onSourceTaskCompleted(mockSrcAttemptId31);
    Assert.assertEquals(4, scheduler.pendingTasks.size()); // no tasks scheduled
    Assert.assertEquals(4, scheduler.numSourceTasks);
    Assert.assertEquals(0, scheduler.numSourceTasksCompleted);
    scheduler.onVertexManagerEventReceived(vmEvent);
    scheduler.onSourceTaskCompleted(mockSrcAttemptId11);
    Assert.assertEquals(4, scheduler.pendingTasks.size());
    Assert.assertEquals(0, scheduledTasks.size()); // no tasks scheduled
    Assert.assertEquals(1, scheduler.numSourceTasksCompleted);
    Assert.assertEquals(1, scheduler.numVertexManagerEventsReceived);
    Assert.assertEquals(500L, scheduler.completedSourceTasksOutputSize);
    // ignore duplicate completion
    scheduler.onSourceTaskCompleted(mockSrcAttemptId11);
    Assert.assertEquals(4, scheduler.pendingTasks.size());
    Assert.assertEquals(0, scheduledTasks.size()); // no tasks scheduled
    Assert.assertEquals(1, scheduler.numSourceTasksCompleted);
    Assert.assertEquals(500L, scheduler.completedSourceTasksOutputSize);

    scheduler.onVertexManagerEventReceived(vmEvent);
    scheduler.onSourceTaskCompleted(mockSrcAttemptId12);
    // managedVertex tasks reduced
    verify(mockManagedVertex).setParallelism(eq(2), anyMap());
    Assert.assertEquals(2, newEdgeManagers.size());
    // TODO improve tests for parallelism
    Assert.assertEquals(0, scheduler.pendingTasks.size()); // all tasks scheduled
    Assert.assertEquals(2, scheduledTasks.size());
    Assert.assertTrue(scheduledTasks.contains(mockTaskId1));
    Assert.assertTrue(scheduledTasks.contains(mockTaskId2));
    Assert.assertEquals(2, scheduler.numSourceTasksCompleted);
    Assert.assertEquals(2, scheduler.numVertexManagerEventsReceived);
    Assert.assertEquals(1000L, scheduler.completedSourceTasksOutputSize);

    // more completions dont cause recalculation of parallelism
    scheduler.onSourceTaskCompleted(mockSrcAttemptId21);
    verify(mockManagedVertex).setParallelism(eq(2), anyMap());
    Assert.assertEquals(2, newEdgeManagers.size());

    EdgeManager edgeManager = newEdgeManagers.values().iterator().next();
    List<Integer> targets = Lists.newArrayList();
    DataMovementEvent dmEvent = new DataMovementEvent(1, new byte[0]);
    edgeManager.routeEventToDestinationTasks(dmEvent, 1, 2, targets);
    Assert.assertEquals(3, dmEvent.getTargetIndex());
    Assert.assertEquals(0, targets.get(0).intValue());
    targets.clear();
    dmEvent = new DataMovementEvent(2, new byte[0]);
    edgeManager.routeEventToDestinationTasks(dmEvent, 0, 2, targets);
    Assert.assertEquals(0, dmEvent.getTargetIndex());
    Assert.assertEquals(1, targets.get(0).intValue());
}

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

License:Apache License

@Test(timeout = 15000l)
public void testDelayedReuseContainerBecomesAvailable()
        throws IOException, InterruptedException, ExecutionException {
    Configuration conf = new Configuration(new YarnConfiguration());
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 3000l);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(conf).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(conf);
    taskSchedulerEventHandler.start();//from w  w w .j a  v a 2  s .com

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();

    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource = Resource.newInstance(1024, 1);
    Priority priority = Priority.newInstance(5);
    String[] host1 = { "host1" };
    String[] host2 = { "host2" };

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

    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 1), 1);
    TezTaskAttemptID taID21 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 2), 1);
    TezTaskAttemptID taID31 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 3), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    TaskAttempt ta21 = mock(TaskAttempt.class);
    TaskAttempt ta31 = mock(TaskAttempt.class);

    AMSchedulerEventTALaunchRequest lrTa11 = createLaunchRequestEvent(taID11, ta11, resource, host1,
            defaultRack, priority);
    AMSchedulerEventTALaunchRequest lrTa21 = createLaunchRequestEvent(taID21, ta21, resource, host2,
            defaultRack, priority);
    AMSchedulerEventTALaunchRequest lrTa31 = createLaunchRequestEvent(taID31, ta31, resource, host1,
            defaultRack, priority);

    taskSchedulerEventHandler.handleEvent(lrTa11);
    taskSchedulerEventHandler.handleEvent(lrTa21);

    Container containerHost1 = createContainer(1, host1[0], resource, priority);
    Container containerHost2 = createContainer(2, host2[0], resource, priority);

    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Lists.newArrayList(containerHost1, containerHost2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta11), any(Object.class), eq(containerHost1));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta21), any(Object.class), eq(containerHost2));

    // Adding the event later so that task1 assigned to containerHost1
    // is deterministic.
    taskSchedulerEventHandler.handleEvent(lrTa31);

    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta11, containerHost1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta11), eq(true));
    verify(taskSchedulerEventHandler, times(1)).taskAllocated(eq(ta31), any(Object.class), eq(containerHost1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(containerHost1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();

    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta21, containerHost2.getId(), TaskAttemptState.SUCCEEDED));

    long currentTs = System.currentTimeMillis();
    Throwable exception = null;
    while (System.currentTimeMillis() < currentTs + 5000l) {
        try {
            verify(taskSchedulerEventHandler, times(1)).containerBeingReleased(eq(containerHost2.getId()));
            exception = null;
            break;
        } catch (Throwable e) {
            exception = e;
        }
    }
    assertTrue("containerHost2 was not released", exception == null);
    taskScheduler.stop();
    taskScheduler.close();
    taskSchedulerEventHandler.close();
}

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

License:Apache License

@Test(timeout = 15000l)
public void testDelayedReuseContainerNotAvailable()
        throws IOException, InterruptedException, ExecutionException {
    Configuration conf = new Configuration(new YarnConfiguration());
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 1000l);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(conf);
    taskSchedulerEventHandler.start();// www .  jav  a  2 s.  c om

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();

    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource = Resource.newInstance(1024, 1);
    Priority priority = Priority.newInstance(5);
    String[] host1 = { "host1" };
    String[] host2 = { "host2" };

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

    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 1), 1);
    TezTaskAttemptID taID21 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 2), 1);
    TezTaskAttemptID taID31 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 3), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    TaskAttempt ta21 = mock(TaskAttempt.class);
    TaskAttempt ta31 = mock(TaskAttempt.class);

    AMSchedulerEventTALaunchRequest lrTa11 = createLaunchRequestEvent(taID11, ta11, resource, host1,
            defaultRack, priority);
    AMSchedulerEventTALaunchRequest lrTa21 = createLaunchRequestEvent(taID21, ta21, resource, host2,
            defaultRack, priority);
    AMSchedulerEventTALaunchRequest lrTa31 = createLaunchRequestEvent(taID31, ta31, resource, host1,
            defaultRack, priority);

    taskSchedulerEventHandler.handleEvent(lrTa11);
    taskSchedulerEventHandler.handleEvent(lrTa21);

    Container containerHost1 = createContainer(1, host1[0], resource, priority);
    Container containerHost2 = createContainer(2, host2[0], resource, priority);

    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Lists.newArrayList(containerHost1, containerHost2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta11), any(Object.class), eq(containerHost1));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta21), any(Object.class), eq(containerHost2));

    // Adding the event later so that task1 assigned to containerHost1 is deterministic.
    taskSchedulerEventHandler.handleEvent(lrTa31);

    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta21, containerHost2.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta21), eq(true));
    verify(taskSchedulerEventHandler, times(0)).taskAllocated(eq(ta31), any(Object.class), eq(containerHost2));
    verify(rmClient, times(1)).releaseAssignedContainer(eq(containerHost2.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);

    taskScheduler.stop();
    taskScheduler.close();
    taskSchedulerEventHandler.close();
}

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

License:Apache License

@Test(timeout = 10000l)
public void testSimpleReuse() throws IOException, InterruptedException, ExecutionException {
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(tezConf);
    taskSchedulerEventHandler.start();/* ww w .j  a  v  a2  s .  c  o  m*/

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource1 = Resource.newInstance(1024, 1);
    String[] host1 = { "host1" };
    String[] host2 = { "host2" };

    String[] racks = { "/default-rack" };
    Priority priority1 = Priority.newInstance(1);

    TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);

    //Vertex 1, Task 1, Attempt 1, host1
    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent1 = createLaunchRequestEvent(taID11, ta11, resource1, host1, racks,
            priority1);

    //Vertex 1, Task 2, Attempt 1, host1
    TezTaskAttemptID taID12 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 2), 1);
    TaskAttempt ta12 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent2 = createLaunchRequestEvent(taID12, ta12, resource1, host1, racks,
            priority1);

    //Vertex 1, Task 3, Attempt 1, host2
    TezTaskAttemptID taID13 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 3), 1);
    TaskAttempt ta13 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent3 = createLaunchRequestEvent(taID13, ta13, resource1, host2, racks,
            priority1);

    //Vertex 1, Task 4, Attempt 1, host2
    TezTaskAttemptID taID14 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 4), 1);
    TaskAttempt ta14 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent4 = createLaunchRequestEvent(taID14, ta14, resource1, host2, racks,
            priority1);

    taskSchedulerEventHandler.handleEvent(lrEvent1);
    taskSchedulerEventHandler.handleEvent(lrEvent2);
    taskSchedulerEventHandler.handleEvent(lrEvent3);
    taskSchedulerEventHandler.handleEvent(lrEvent4);

    Container container1 = createContainer(1, "host1", resource1, priority1);

    // One container allocated.
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta11), any(Object.class), eq(container1));

    // Task assigned to container completed successfully. Container should be re-used.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta11), eq(true));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta12), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();

    // Task assigned to container completed successfully.
    // Verify reuse across hosts.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta12, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta12), eq(true));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta13), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();

    // Verify no re-use if a previous task fails.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta13, container1.getId(), TaskAttemptState.FAILED));
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler, times(0)).taskAllocated(eq(ta14), any(Object.class), eq(container1));
    verify(taskScheduler).deallocateTask(eq(ta13), eq(false));
    verify(rmClient).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();

    Container container2 = createContainer(2, "host2", resource1, priority1);

    // Second container allocated. Should be allocated to the last task.
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta14), any(Object.class), eq(container2));

    // Task assigned to container completed successfully. No pending requests. Container should be released.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta14, container2.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta14), eq(true));
    verify(rmClient).releaseAssignedContainer(eq(container2.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();

    taskScheduler.close();
    taskSchedulerEventHandler.close();
}

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

License:Apache License

@Test(timeout = 10000l)
public void testReuseWithTaskSpecificLaunchCmdOption()
        throws IOException, InterruptedException, ExecutionException {
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
    //Profile 3 tasks
    tezConf.set(TezConfiguration.TEZ_TASK_SPECIFIC_LAUNCH_CMD_OPTS_LIST, "v1[1,3,4]");
    tezConf.set(TezConfiguration.TEZ_TASK_SPECIFIC_LAUNCH_CMD_OPTS, "dir=/tmp/__VERTEX_NAME__/__TASK_INDEX__");
    TaskSpecificLaunchCmdOption taskSpecificLaunchCmdOption = new TaskSpecificLaunchCmdOption(tezConf);

    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    //Use ContainerContextMatcher here.  Otherwise it would not match the JVM options
    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new ContainerContextMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(tezConf);
    taskSchedulerEventHandler.start();//  ww w  . jav a  2s. co  m

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource1 = Resource.newInstance(1024, 1);
    String[] host1 = { "host1" };
    String[] host2 = { "host2" };
    String[] host3 = { "host3" };

    String[] racks = { "/default-rack" };
    Priority priority1 = Priority.newInstance(1);

    TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);
    Map<String, LocalResource> localResources = new HashMap<String, LocalResource>();
    String tsLaunchCmdOpts = taskSpecificLaunchCmdOption.getTaskSpecificOption("", "v1", 1);

    /**
     * Schedule 2 tasks (1 with additional launch-cmd option and another in normal mode).
     * Container should not be reused in this case.
     */
    //Vertex 1, Task 1, Attempt 1, host1
    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent1 = createLaunchRequestEvent(taID11, ta11, resource1, host1, racks,
            priority1, localResources, tsLaunchCmdOpts);

    //Vertex 1, Task 2, Attempt 1, host1
    TezTaskAttemptID taID12 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 2), 1);
    TaskAttempt ta12 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent2 = createLaunchRequestEvent(taID12, ta12, resource1, host1, racks,
            priority1);

    taskSchedulerEventHandler.handleEvent(lrEvent1);
    taskSchedulerEventHandler.handleEvent(lrEvent2);

    Container container1 = createContainer(1, "host1", resource1, priority1);

    // One container allocated.
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta11), any(Object.class), eq(container1));

    // First task had profiling on. This container can not be reused further.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta11), eq(true));
    verify(taskSchedulerEventHandler, times(0)).taskAllocated(eq(ta12), any(Object.class), eq(container1));
    verify(rmClient, times(1)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();

    /**
     * Schedule 2 tasks (both having different task specific JVM option).
     * Container should not be reused.
     */
    //Vertex 1, Task 3, Attempt 1, host2
    tsLaunchCmdOpts = taskSpecificLaunchCmdOption.getTaskSpecificOption("", "v1", 3);
    TezTaskAttemptID taID13 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 3), 1);
    TaskAttempt ta13 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent3 = createLaunchRequestEvent(taID13, ta13, resource1, host2, racks,
            priority1, localResources, tsLaunchCmdOpts);

    //Vertex 1, Task 4, Attempt 1, host2
    tsLaunchCmdOpts = taskSpecificLaunchCmdOption.getTaskSpecificOption("", "v1", 4);
    TezTaskAttemptID taID14 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 4), 1);
    TaskAttempt ta14 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent4 = createLaunchRequestEvent(taID14, ta14, resource1, host2, racks,
            priority1, localResources, tsLaunchCmdOpts);

    Container container2 = createContainer(2, "host2", resource1, priority1);
    taskSchedulerEventHandler.handleEvent(lrEvent3);
    taskSchedulerEventHandler.handleEvent(lrEvent4);

    // Container started
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta13), any(Object.class), eq(container2));

    // Verify that the container can not be reused when profiling option is turned on
    // Even for 2 tasks having same profiling option can have container reusability.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta13, container2.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta13), eq(true));
    verify(taskSchedulerEventHandler, times(0)).taskAllocated(eq(ta14), any(Object.class), eq(container2));
    verify(rmClient, times(1)).releaseAssignedContainer(eq(container2.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();

    /**
     * Schedule 2 tasks with same jvm profiling option.
     * Container should be reused.
     */
    tezConf.set(TezConfiguration.TEZ_TASK_SPECIFIC_LAUNCH_CMD_OPTS_LIST, "v1[1,2,3,5,6]");
    tezConf.set(TezConfiguration.TEZ_TASK_SPECIFIC_LAUNCH_CMD_OPTS, "dummyOpts");
    taskSpecificLaunchCmdOption = new TaskSpecificLaunchCmdOption(tezConf);

    //Vertex 1, Task 5, Attempt 1, host3
    TezTaskAttemptID taID15 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 3), 1);
    TaskAttempt ta15 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent5 = createLaunchRequestEvent(taID15, ta15, resource1, host3, racks,
            priority1, localResources, taskSpecificLaunchCmdOption.getTaskSpecificOption("", "v1", 5));

    //Vertex 1, Task 6, Attempt 1, host3
    tsLaunchCmdOpts = taskSpecificLaunchCmdOption.getTaskSpecificOption("", "v1", 4);
    TezTaskAttemptID taID16 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 4), 1);
    TaskAttempt ta16 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent6 = createLaunchRequestEvent(taID16, ta16, resource1, host3, racks,
            priority1, localResources, taskSpecificLaunchCmdOption.getTaskSpecificOption("", "v1", 6));

    // Container started
    Container container3 = createContainer(2, "host3", resource1, priority1);
    taskSchedulerEventHandler.handleEvent(lrEvent5);
    taskSchedulerEventHandler.handleEvent(lrEvent6);

    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container3));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta15), any(Object.class), eq(container3));

    //Ensure task 6 (of vertex 1) is allocated to same container
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta15, container3.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta15), eq(true));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta16), any(Object.class), eq(container3));
    eventHandler.reset();

    taskScheduler.close();
    taskSchedulerEventHandler.close();
}

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

License:Apache License

@Test(timeout = 30000l)
public void testReuseNonLocalRequest() throws IOException, InterruptedException, ExecutionException {
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100l);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 1000l);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 1000l);

    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(tezConf);
    taskSchedulerEventHandler.start();//from w  w w.  j  av a2s.  co  m

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource1 = Resource.newInstance(1024, 1);
    String[] emptyHosts = new String[0];
    String[] racks = { "default-rack" };

    Priority priority = Priority.newInstance(3);

    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);

    //Vertex 1, Task 1, Attempt 1, no locality information.
    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 1), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    doReturn(vertexID).when(ta11).getVertexID();
    AMSchedulerEventTALaunchRequest lrEvent11 = createLaunchRequestEvent(taID11, ta11, resource1, emptyHosts,
            racks, priority);

    //Vertex1, Task2, Attempt 1,  no locality information.
    TezTaskAttemptID taID12 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID, 2), 1);
    TaskAttempt ta12 = mock(TaskAttempt.class);
    doReturn(vertexID).when(ta12).getVertexID();
    AMSchedulerEventTALaunchRequest lrEvent12 = createLaunchRequestEvent(taID12, ta12, resource1, emptyHosts,
            racks, priority);

    // Send launch request for task 1 only, deterministic assignment to this task.
    taskSchedulerEventHandler.handleEvent(lrEvent11);

    Container container1 = createContainer(1, "randomHost", resource1, priority);

    // One container allocated.
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta11), any(Object.class), eq(container1));

    // Send launch request for task2 (vertex2)
    taskSchedulerEventHandler.handleEvent(lrEvent12);

    // Task assigned to container completed successfully.
    // Container should not be immediately assigned to task 2
    // until delay expires.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta11), eq(true));
    verify(taskSchedulerEventHandler, times(0)).taskAllocated(eq(ta12), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();

    LOG.info("Sleeping to ensure that the scheduling loop runs");
    Thread.sleep(3000l);
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta12), any(Object.class), eq(container1));

    // TA12 completed.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta12, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    LOG.info("Sleeping to ensure that the scheduling loop runs");
    Thread.sleep(3000l);
    verify(rmClient).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);

    taskScheduler.close();
    taskSchedulerEventHandler.close();
}

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

License:Apache License

@Test(timeout = 30000l)
public void testReuseAcrossVertices() throws IOException, InterruptedException, ExecutionException {
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 1l);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 2000l);
    tezConf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 1);

    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(true).when(appContext).isSession();
    doReturn(dagID).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(tezConf);
    taskSchedulerEventHandler.start();/*w  w w  .  ja  va  2  s .  c o m*/

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();

    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource1 = Resource.newInstance(1024, 1);
    String[] host1 = { "host1" };

    String[] racks = { "/default-rack" };
    Priority priority1 = Priority.newInstance(3);
    Priority priority2 = Priority.newInstance(4);

    TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);
    TezVertexID vertexID2 = TezVertexID.getInstance(dagID, 2);

    //Vertex 1, Task 1, Attempt 1, host1
    TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
    TaskAttempt ta11 = mock(TaskAttempt.class);
    doReturn(vertexID1).when(ta11).getVertexID();
    AMSchedulerEventTALaunchRequest lrEvent11 = createLaunchRequestEvent(taID11, ta11, resource1, host1, racks,
            priority1);

    //Vertex2, Task1, Attempt 1, host1
    TezTaskAttemptID taID21 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID2, 1), 1);
    TaskAttempt ta21 = mock(TaskAttempt.class);
    doReturn(vertexID2).when(ta21).getVertexID();
    AMSchedulerEventTALaunchRequest lrEvent21 = createLaunchRequestEvent(taID21, ta21, resource1, host1, racks,
            priority2);

    // Send launch request for task 1 onle, deterministic assignment to this task.
    taskSchedulerEventHandler.handleEvent(lrEvent11);

    Container container1 = createContainer(1, host1[0], resource1, priority1);

    // One container allocated.
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta11), any(Object.class), eq(container1));

    // Send launch request for task2 (vertex2)
    taskSchedulerEventHandler.handleEvent(lrEvent21);

    // Task assigned to container completed successfully.
    // Container should  be assigned to task21.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta11), eq(true));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta21), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));

    // Task 2 completes.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta21, container1.getId(), TaskAttemptState.SUCCEEDED));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));

    LOG.info("Sleeping to ensure that the scheduling loop runs");
    Thread.sleep(3000l);
    // container should not get released due to min held containers
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));

    taskScheduler.close();
    taskSchedulerEventHandler.close();
}

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

License:Apache License

@Test(timeout = 30000l)
public void testReuseLocalResourcesChanged() throws IOException, InterruptedException, ExecutionException {
    Configuration tezConf = new Configuration(new YarnConfiguration());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, true);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, -1);

    TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class);

    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID1 = TezDAGID.getInstance("0", 1, 0);

    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    String appUrl = "url";
    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);

    doReturn(finalStatus).when(mockApp).getFinalAppStatus();

    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    ChangingDAGIDAnswer dagIDAnswer = new ChangingDAGIDAnswer(dagID1);
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appContext);
    AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    doReturn(amNodeTracker).when(appContext).getNodeTracker();
    doReturn(DAGAppMasterState.RUNNING).when(appContext).getAMState();
    doReturn(true).when(appContext).isSession();
    doAnswer(dagIDAnswer).when(appContext).getCurrentDAGID();
    doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

    TaskSchedulerEventHandler taskSchedulerEventHandlerReal = new TaskSchedulerEventHandlerForTest(appContext,
            eventHandler, rmClient, new AlwaysMatchesContainerMatcher());
    TaskSchedulerEventHandler taskSchedulerEventHandler = spy(taskSchedulerEventHandlerReal);
    taskSchedulerEventHandler.init(tezConf);
    taskSchedulerEventHandler.start();/*  w  w w  .  jav  a  2 s.c  om*/

    TaskSchedulerWithDrainableAppCallback taskScheduler = (TaskSchedulerWithDrainableAppCallback) ((TaskSchedulerEventHandlerForTest) taskSchedulerEventHandler)
            .getSpyTaskScheduler();
    TaskSchedulerAppCallbackDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
    AtomicBoolean drainNotifier = new AtomicBoolean(false);
    taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

    Resource resource1 = Resource.newInstance(1024, 1);
    String[] host1 = { "host1" };

    String[] racks = { "/default-rack" };
    Priority priority1 = Priority.newInstance(1);

    String rsrc1 = "rsrc1";
    String rsrc2 = "rsrc2";
    String rsrc3 = "rsrc3";
    LocalResource lr1 = mock(LocalResource.class);
    LocalResource lr2 = mock(LocalResource.class);
    LocalResource lr3 = mock(LocalResource.class);

    AMContainerEventAssignTA assignEvent = null;

    Map<String, LocalResource> dag1LRs = Maps.newHashMap();
    dag1LRs.put(rsrc1, lr1);

    TezVertexID vertexID11 = TezVertexID.getInstance(dagID1, 1);

    //Vertex 1, Task 1, Attempt 1, host1, lr1
    TezTaskAttemptID taID111 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID11, 1), 1);
    TaskAttempt ta111 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent11 = createLaunchRequestEvent(taID111, ta111, resource1, host1,
            racks, priority1, dag1LRs);

    //Vertex 1, Task 2, Attempt 1, host1, lr1
    TezTaskAttemptID taID112 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID11, 2), 1);
    TaskAttempt ta112 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent12 = createLaunchRequestEvent(taID112, ta112, resource1, host1,
            racks, priority1, dag1LRs);

    taskSchedulerEventHandler.handleEvent(lrEvent11);
    taskSchedulerEventHandler.handleEvent(lrEvent12);

    Container container1 = createContainer(1, "host1", resource1, priority1);

    // One container allocated.
    drainNotifier.set(false);
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta111), any(Object.class), eq(container1));
    assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
    assertEquals(1, assignEvent.getRemoteTaskLocalResources().size());

    // Task assigned to container completed successfully. Container should be re-used.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta111, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta111), eq(true));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta112), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
    assertEquals(1, assignEvent.getRemoteTaskLocalResources().size());
    eventHandler.reset();

    // Task assigned to container completed successfully.
    // Verify reuse across hosts.
    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta112, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta112), eq(true));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();

    // Setup DAG2 with additional resources. Make sure the container, even without all resources, is reused.
    TezDAGID dagID2 = TezDAGID.getInstance("0", 2, 0);
    dagIDAnswer.setDAGID(dagID2);

    Map<String, LocalResource> dag2LRs = Maps.newHashMap();
    dag2LRs.put(rsrc2, lr2);
    dag2LRs.put(rsrc3, lr3);

    TezVertexID vertexID21 = TezVertexID.getInstance(dagID2, 1);

    //Vertex 2, Task 1, Attempt 1, host1, lr2
    TezTaskAttemptID taID211 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID21, 1), 1);
    TaskAttempt ta211 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent21 = createLaunchRequestEvent(taID211, ta211, resource1, host1,
            racks, priority1, dag2LRs);

    //Vertex 2, Task 2, Attempt 1, host1, lr2
    TezTaskAttemptID taID212 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID21, 2), 1);
    TaskAttempt ta212 = mock(TaskAttempt.class);
    AMSchedulerEventTALaunchRequest lrEvent22 = createLaunchRequestEvent(taID212, ta212, resource1, host1,
            racks, priority1, dag2LRs);

    taskSchedulerEventHandler.handleEvent(lrEvent21);
    taskSchedulerEventHandler.handleEvent(lrEvent22);
    drainableAppCallback.drain();

    // TODO This is terrible, need a better way to ensure the scheduling loop has run
    LOG.info("Sleeping to ensure that the scheduling loop runs");
    Thread.sleep(6000l);
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta211), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
    assertEquals(2, assignEvent.getRemoteTaskLocalResources().size());
    eventHandler.reset();

    taskSchedulerEventHandler
            .handleEvent(new AMSchedulerEventTAEnded(ta211, container1.getId(), TaskAttemptState.SUCCEEDED));
    drainableAppCallback.drain();
    verify(taskScheduler).deallocateTask(eq(ta211), eq(true));
    verify(taskSchedulerEventHandler).taskAllocated(eq(ta212), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    assignEvent = (AMContainerEventAssignTA) eventHandler.verifyInvocation(AMContainerEventAssignTA.class);
    assertEquals(2, assignEvent.getRemoteTaskLocalResources().size());
    eventHandler.reset();

    taskScheduler.close();
    taskSchedulerEventHandler.close();
}