Search in sources :

Example 91 with TezVertexID

use of org.apache.tez.dag.records.TezVertexID in project tez by apache.

the class TestEdge method testCompositeEventHandling.

@SuppressWarnings({ "rawtypes" })
@Test(timeout = 5000)
public void testCompositeEventHandling() throws TezException {
    EventHandler eventHandler = mock(EventHandler.class);
    EdgeProperty edgeProp = EdgeProperty.create(DataMovementType.SCATTER_GATHER, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, mock(OutputDescriptor.class), mock(InputDescriptor.class));
    Edge edge = new Edge(edgeProp, eventHandler, new TezConfiguration());
    TezVertexID srcVertexID = createVertexID(1);
    TezVertexID destVertexID = createVertexID(2);
    LinkedHashMap<TezTaskID, Task> srcTasks = mockTasks(srcVertexID, 1);
    LinkedHashMap<TezTaskID, Task> destTasks = mockTasks(destVertexID, 5);
    TezTaskID srcTaskID = srcTasks.keySet().iterator().next();
    Vertex srcVertex = mockVertex("src", srcVertexID, srcTasks);
    Vertex destVertex = mockVertex("dest", destVertexID, destTasks);
    edge.setSourceVertex(srcVertex);
    edge.setDestinationVertex(destVertex);
    edge.initialize();
    // Task0, Attempt 0
    TezTaskAttemptID srcTAID = createTAIDForTest(srcTaskID, 2);
    EventMetaData srcMeta = new EventMetaData(EventProducerConsumerType.OUTPUT, "consumerVertex", "producerVertex", srcTAID);
    // Verification via a CompositeEvent
    CompositeDataMovementEvent cdmEvent = CompositeDataMovementEvent.create(0, destTasks.size(), ByteBuffer.wrap("bytes".getBytes()));
    // AttemptNum
    cdmEvent.setVersion(2);
    TezEvent tezEvent = new TezEvent(cdmEvent, srcMeta);
    // Event setup to look like it would after the Vertex is done with it.
    edge.sendTezEventToDestinationTasks(tezEvent);
    verifyEvents(srcTAID, destTasks);
    // Same Verification via regular DataMovementEvents
    // Reset the mock
    resetTaskMocks(destTasks.values());
    for (int i = 0; i < destTasks.size(); i++) {
        DataMovementEvent dmEvent = DataMovementEvent.create(i, ByteBuffer.wrap("bytes".getBytes()));
        dmEvent.setVersion(2);
        tezEvent = new TezEvent(dmEvent, srcMeta);
        edge.sendTezEventToDestinationTasks(tezEvent);
    }
    verifyEvents(srcTAID, destTasks);
}
Also used : InputDescriptor(org.apache.tez.dag.api.InputDescriptor) Vertex(org.apache.tez.dag.app.dag.Vertex) Task(org.apache.tez.dag.app.dag.Task) EventHandler(org.apache.hadoop.yarn.event.EventHandler) DataMovementEvent(org.apache.tez.runtime.api.events.DataMovementEvent) CompositeDataMovementEvent(org.apache.tez.runtime.api.events.CompositeDataMovementEvent) TezTaskID(org.apache.tez.dag.records.TezTaskID) OutputDescriptor(org.apache.tez.dag.api.OutputDescriptor) CompositeDataMovementEvent(org.apache.tez.runtime.api.events.CompositeDataMovementEvent) EdgeProperty(org.apache.tez.dag.api.EdgeProperty) TezEvent(org.apache.tez.runtime.api.impl.TezEvent) TezVertexID(org.apache.tez.dag.records.TezVertexID) EventMetaData(org.apache.tez.runtime.api.impl.EventMetaData) TezConfiguration(org.apache.tez.dag.api.TezConfiguration) TezTaskAttemptID(org.apache.tez.dag.records.TezTaskAttemptID) Test(org.junit.Test) EdgeManagerForTest(org.apache.tez.test.EdgeManagerForTest)

Example 92 with TezVertexID

use of org.apache.tez.dag.records.TezVertexID in project tez by apache.

the class TestEdge method createVertexID.

private TezVertexID createVertexID(int id) {
    TezDAGID dagID = TezDAGID.getInstance("1000", 1, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, id);
    return vertexID;
}
Also used : TezDAGID(org.apache.tez.dag.records.TezDAGID) TezVertexID(org.apache.tez.dag.records.TezVertexID)

Example 93 with TezVertexID

use of org.apache.tez.dag.records.TezVertexID in project tez by apache.

the class TestContainerReuse method testDelayedReuseContainerBecomesAvailable.

@Test(timeout = 15000l)
public void testDelayedReuseContainerBecomesAvailable() throws IOException, InterruptedException, ExecutionException {
    LOG.info("Test testDelayedReuseContainerBecomesAvailable");
    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);
    conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);
    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));
    AppContext appContext = mock(AppContext.class);
    doReturn(conf).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.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();
    TaskSchedulerManager taskSchedulerManagerReal = new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(conf));
    TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
    taskSchedulerManager.init(conf);
    taskSchedulerManager.start();
    TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext) ((TaskSchedulerManagerForTest) taskSchedulerManager).getSpyTaskScheduler();
    TaskSchedulerContextDrainable 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);
    taskSchedulerManager.handleEvent(lrTa11);
    taskSchedulerManager.handleEvent(lrTa21);
    Container containerHost1 = createContainer(1, host1[0], resource, priority);
    Container containerHost2 = createContainer(2, host2[0], resource, priority);
    taskScheduler.onContainersAllocated(Lists.newArrayList(containerHost1, containerHost2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta11), any(Object.class), eq(containerHost1));
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta21), any(Object.class), eq(containerHost2));
    // Adding the event later so that task1 assigned to containerHost1
    // is deterministic.
    taskSchedulerManager.handleEvent(lrTa31);
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta11, containerHost1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta11, true, null, null);
    verify(taskSchedulerManager, times(1)).taskAllocated(eq(0), eq(ta31), any(Object.class), eq(containerHost1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(containerHost1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta21, containerHost2.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    long currentTs = System.currentTimeMillis();
    Throwable exception = null;
    while (System.currentTimeMillis() < currentTs + 5000l) {
        try {
            verify(taskSchedulerManager, times(1)).containerBeingReleased(eq(0), eq(containerHost2.getId()));
            exception = null;
            break;
        } catch (Throwable e) {
            exception = e;
        }
    }
    assertTrue("containerHost2 was not released", exception == null);
    taskScheduler.shutdown();
    taskSchedulerManager.close();
}
Also used : CapturingEventHandler(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.CapturingEventHandler) Configuration(org.apache.hadoop.conf.Configuration) TezConfiguration(org.apache.tez.dag.api.TezConfiguration) YarnConfiguration(org.apache.hadoop.yarn.conf.YarnConfiguration) TaskSchedulerWithDrainableContext(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext) AMContainerMap(org.apache.tez.dag.app.rm.container.AMContainerMap) AlwaysMatchesContainerMatcher(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AlwaysMatchesContainerMatcher) Container(org.apache.hadoop.yarn.api.records.Container) YarnConfiguration(org.apache.hadoop.yarn.conf.YarnConfiguration) TezDAGID(org.apache.tez.dag.records.TezDAGID) AMRMClientAsyncForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientAsyncForTest) TaskAttempt(org.apache.tez.dag.app.dag.TaskAttempt) TezVertexID(org.apache.tez.dag.records.TezVertexID) ContainerHeartbeatHandler(org.apache.tez.dag.app.ContainerHeartbeatHandler) TaskSchedulerContextDrainable(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerContextDrainable) Priority(org.apache.hadoop.yarn.api.records.Priority) AppContext(org.apache.tez.dag.app.AppContext) Resource(org.apache.hadoop.yarn.api.records.Resource) LocalResource(org.apache.hadoop.yarn.api.records.LocalResource) TaskCommunicatorManagerInterface(org.apache.tez.dag.app.TaskCommunicatorManagerInterface) ContainerContextMatcher(org.apache.tez.dag.app.rm.container.ContainerContextMatcher) AMNodeTracker(org.apache.tez.dag.app.rm.node.AMNodeTracker) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ClusterInfo(org.apache.tez.dag.app.ClusterInfo) TaskSchedulerManagerForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerManagerForTest) CookieContainerRequest(org.apache.tez.dag.app.rm.YarnTaskSchedulerService.CookieContainerRequest) AMRMClientForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientForTest) TezTaskAttemptID(org.apache.tez.dag.records.TezTaskAttemptID) TaskSchedulerManagerForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerManagerForTest) AMRMClientForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientForTest) AMRMClientAsyncForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientAsyncForTest) Test(org.junit.Test)

Example 94 with TezVertexID

use of org.apache.tez.dag.records.TezVertexID in project tez by apache.

the class TestContainerReuse method testReuseNonLocalRequest.

@Test(timeout = 30000l)
public void testReuseNonLocalRequest() throws IOException, InterruptedException, ExecutionException {
    LOG.info("Test testReuseNonLocalRequest");
    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);
    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.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();
    TaskSchedulerManager taskSchedulerManagerReal = new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(tezConf));
    TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
    taskSchedulerManager.init(tezConf);
    taskSchedulerManager.start();
    TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext) ((TaskSchedulerManagerForTest) taskSchedulerManager).getSpyTaskScheduler();
    TaskSchedulerContextDrainable 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.
    taskSchedulerManager.handleEvent(lrEvent11);
    Container container1 = createContainer(1, "randomHost", resource1, priority);
    // One container allocated.
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta11), any(Object.class), eq(container1));
    // Send launch request for task2 (vertex2)
    taskSchedulerManager.handleEvent(lrEvent12);
    // Task assigned to container completed successfully.
    // Container should not be immediately assigned to task 2
    // until delay expires.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta11, true, null, null);
    verify(taskSchedulerManager, times(0)).taskAllocated(eq(0), eq(ta12), any(Object.class), eq(container1));
    verify(rmClient, times(0)).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyNoInvocations(AMContainerEventStopRequest.class);
    eventHandler.reset();
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta12), any(Object.class), eq(container1));
    // TA12 completed.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta12, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(rmClient).releaseAssignedContainer(eq(container1.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    taskScheduler.shutdown();
    taskSchedulerManager.close();
}
Also used : CapturingEventHandler(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.CapturingEventHandler) Configuration(org.apache.hadoop.conf.Configuration) TezConfiguration(org.apache.tez.dag.api.TezConfiguration) YarnConfiguration(org.apache.hadoop.yarn.conf.YarnConfiguration) TaskSchedulerWithDrainableContext(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext) AMContainerMap(org.apache.tez.dag.app.rm.container.AMContainerMap) AlwaysMatchesContainerMatcher(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AlwaysMatchesContainerMatcher) Container(org.apache.hadoop.yarn.api.records.Container) YarnConfiguration(org.apache.hadoop.yarn.conf.YarnConfiguration) TezDAGID(org.apache.tez.dag.records.TezDAGID) AMRMClientAsyncForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientAsyncForTest) TaskAttempt(org.apache.tez.dag.app.dag.TaskAttempt) ContainerHeartbeatHandler(org.apache.tez.dag.app.ContainerHeartbeatHandler) TaskSchedulerContextDrainable(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerContextDrainable) TezVertexID(org.apache.tez.dag.records.TezVertexID) Priority(org.apache.hadoop.yarn.api.records.Priority) AppContext(org.apache.tez.dag.app.AppContext) Resource(org.apache.hadoop.yarn.api.records.Resource) LocalResource(org.apache.hadoop.yarn.api.records.LocalResource) TaskCommunicatorManagerInterface(org.apache.tez.dag.app.TaskCommunicatorManagerInterface) ContainerContextMatcher(org.apache.tez.dag.app.rm.container.ContainerContextMatcher) AMNodeTracker(org.apache.tez.dag.app.rm.node.AMNodeTracker) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ClusterInfo(org.apache.tez.dag.app.ClusterInfo) TaskSchedulerManagerForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerManagerForTest) CookieContainerRequest(org.apache.tez.dag.app.rm.YarnTaskSchedulerService.CookieContainerRequest) AMRMClientForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientForTest) TezTaskAttemptID(org.apache.tez.dag.records.TezTaskAttemptID) TaskSchedulerManagerForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerManagerForTest) AMRMClientForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientForTest) AMRMClientAsyncForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientAsyncForTest) Test(org.junit.Test)

Example 95 with TezVertexID

use of org.apache.tez.dag.records.TezVertexID in project tez by apache.

the class TestContainerReuse method testSimpleReuse.

@Test(timeout = 10000l)
public void testSimpleReuse() throws IOException, InterruptedException, ExecutionException {
    LOG.info("Test testSimpleReuse");
    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);
    tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);
    CapturingEventHandler eventHandler = new CapturingEventHandler();
    TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);
    AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
    TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));
    AppContext appContext = mock(AppContext.class);
    doReturn(new Configuration(false)).when(appContext).getAMConf();
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), mock(TaskCommunicatorManagerInterface.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();
    TaskSchedulerManager taskSchedulerManagerReal = new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient, new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(tezConf));
    TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
    taskSchedulerManager.init(tezConf);
    taskSchedulerManager.start();
    TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext) ((TaskSchedulerManagerForTest) taskSchedulerManager).getSpyTaskScheduler();
    TaskSchedulerContextDrainable 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);
    taskSchedulerManager.handleEvent(lrEvent1);
    taskSchedulerManager.handleEvent(lrEvent2);
    taskSchedulerManager.handleEvent(lrEvent3);
    taskSchedulerManager.handleEvent(lrEvent4);
    Container container1 = createContainer(1, "host1", resource1, priority1);
    // One container allocated.
    taskScheduler.onContainersAllocated(Collections.singletonList(container1));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta11), any(Object.class), eq(container1));
    // Task assigned to container completed successfully. Container should be re-used.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta11, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta11, true, null, null);
    verify(taskSchedulerManager).taskAllocated(eq(0), 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.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta12, container1.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta12, true, null, null);
    verify(taskSchedulerManager).taskAllocated(eq(0), 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.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta13, container1.getId(), TaskAttemptState.FAILED, null, "TIMEOUT", 0));
    drainableAppCallback.drain();
    verify(taskSchedulerManager, times(0)).taskAllocated(eq(0), eq(ta14), any(Object.class), eq(container1));
    verifyDeAllocateTask(taskScheduler, ta13, false, null, "TIMEOUT");
    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.
    taskScheduler.onContainersAllocated(Collections.singletonList(container2));
    TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier);
    drainableAppCallback.drain();
    verify(taskSchedulerManager).taskAllocated(eq(0), eq(ta14), any(Object.class), eq(container2));
    // Task assigned to container completed successfully. No pending requests. Container should be released.
    taskSchedulerManager.handleEvent(new AMSchedulerEventTAEnded(ta14, container2.getId(), TaskAttemptState.SUCCEEDED, null, null, 0));
    drainableAppCallback.drain();
    verifyDeAllocateTask(taskScheduler, ta14, true, null, null);
    verify(rmClient).releaseAssignedContainer(eq(container2.getId()));
    eventHandler.verifyInvocation(AMContainerEventStopRequest.class);
    eventHandler.reset();
    taskScheduler.shutdown();
    taskSchedulerManager.close();
}
Also used : CapturingEventHandler(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.CapturingEventHandler) Configuration(org.apache.hadoop.conf.Configuration) TezConfiguration(org.apache.tez.dag.api.TezConfiguration) YarnConfiguration(org.apache.hadoop.yarn.conf.YarnConfiguration) TaskSchedulerWithDrainableContext(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerWithDrainableContext) AMContainerMap(org.apache.tez.dag.app.rm.container.AMContainerMap) AlwaysMatchesContainerMatcher(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AlwaysMatchesContainerMatcher) Container(org.apache.hadoop.yarn.api.records.Container) YarnConfiguration(org.apache.hadoop.yarn.conf.YarnConfiguration) TezDAGID(org.apache.tez.dag.records.TezDAGID) AMRMClientAsyncForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientAsyncForTest) TaskAttempt(org.apache.tez.dag.app.dag.TaskAttempt) ContainerHeartbeatHandler(org.apache.tez.dag.app.ContainerHeartbeatHandler) TaskSchedulerContextDrainable(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerContextDrainable) TezVertexID(org.apache.tez.dag.records.TezVertexID) Priority(org.apache.hadoop.yarn.api.records.Priority) AppContext(org.apache.tez.dag.app.AppContext) Resource(org.apache.hadoop.yarn.api.records.Resource) LocalResource(org.apache.hadoop.yarn.api.records.LocalResource) TaskCommunicatorManagerInterface(org.apache.tez.dag.app.TaskCommunicatorManagerInterface) ContainerContextMatcher(org.apache.tez.dag.app.rm.container.ContainerContextMatcher) AMNodeTracker(org.apache.tez.dag.app.rm.node.AMNodeTracker) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ClusterInfo(org.apache.tez.dag.app.ClusterInfo) TaskSchedulerManagerForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerManagerForTest) CookieContainerRequest(org.apache.tez.dag.app.rm.YarnTaskSchedulerService.CookieContainerRequest) AMRMClientForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientForTest) TezTaskAttemptID(org.apache.tez.dag.records.TezTaskAttemptID) TaskSchedulerManagerForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.TaskSchedulerManagerForTest) AMRMClientForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientForTest) AMRMClientAsyncForTest(org.apache.tez.dag.app.rm.TestTaskSchedulerHelpers.AMRMClientAsyncForTest) Test(org.junit.Test)

Aggregations

TezVertexID (org.apache.tez.dag.records.TezVertexID)117 Test (org.junit.Test)64 TezDAGID (org.apache.tez.dag.records.TezDAGID)61 TezTaskAttemptID (org.apache.tez.dag.records.TezTaskAttemptID)54 TezConfiguration (org.apache.tez.dag.api.TezConfiguration)46 Configuration (org.apache.hadoop.conf.Configuration)43 TezTaskID (org.apache.tez.dag.records.TezTaskID)43 ApplicationId (org.apache.hadoop.yarn.api.records.ApplicationId)40 LocalResource (org.apache.hadoop.yarn.api.records.LocalResource)31 Resource (org.apache.hadoop.yarn.api.records.Resource)29 Container (org.apache.hadoop.yarn.api.records.Container)28 ClusterInfo (org.apache.tez.dag.app.ClusterInfo)28 TaskCommunicatorManagerInterface (org.apache.tez.dag.app.TaskCommunicatorManagerInterface)28 ContainerHeartbeatHandler (org.apache.tez.dag.app.ContainerHeartbeatHandler)27 AMContainerMap (org.apache.tez.dag.app.rm.container.AMContainerMap)27 ContainerContextMatcher (org.apache.tez.dag.app.rm.container.ContainerContextMatcher)27 ContainerId (org.apache.hadoop.yarn.api.records.ContainerId)24 TezEvent (org.apache.tez.runtime.api.impl.TezEvent)24 SystemClock (org.apache.hadoop.yarn.util.SystemClock)22 Vertex (org.apache.tez.dag.app.dag.Vertex)22