| /** |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| package org.apache.tez.dag.app.rm; |
| |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertTrue; |
| import static org.junit.Assert.fail; |
| import static org.mockito.Matchers.any; |
| import static org.mockito.Matchers.eq; |
| import static org.mockito.Mockito.RETURNS_DEEP_STUBS; |
| import static org.mockito.Mockito.doReturn; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.reset; |
| import static org.mockito.Mockito.spy; |
| import static org.mockito.Mockito.times; |
| import static org.mockito.Mockito.verify; |
| import static org.mockito.Mockito.when; |
| |
| import javax.annotation.Nullable; |
| import java.io.IOException; |
| import java.lang.reflect.Method; |
| import java.net.InetSocketAddress; |
| import java.nio.ByteBuffer; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.LinkedList; |
| import java.util.List; |
| import java.util.Set; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.atomic.AtomicBoolean; |
| import java.util.concurrent.atomic.AtomicInteger; |
| |
| import org.apache.hadoop.conf.Configuration; |
| import org.apache.hadoop.security.Credentials; |
| import org.apache.hadoop.yarn.api.records.ApplicationId; |
| import org.apache.hadoop.yarn.api.records.Container; |
| import org.apache.hadoop.yarn.api.records.ContainerExitStatus; |
| import org.apache.hadoop.yarn.api.records.ContainerId; |
| import org.apache.hadoop.yarn.api.records.ContainerStatus; |
| import org.apache.hadoop.yarn.api.records.LocalResource; |
| import org.apache.hadoop.yarn.api.records.NodeId; |
| import org.apache.hadoop.yarn.api.records.Priority; |
| import org.apache.hadoop.yarn.api.records.Resource; |
| import org.apache.hadoop.yarn.event.Event; |
| import org.apache.hadoop.yarn.event.EventHandler; |
| import org.apache.tez.common.ContainerSignatureMatcher; |
| import org.apache.tez.common.TezUtils; |
| import org.apache.tez.dag.api.NamedEntityDescriptor; |
| import org.apache.tez.dag.api.TaskLocationHint; |
| import org.apache.tez.dag.api.TezConfiguration; |
| import org.apache.tez.dag.api.TezConstants; |
| import org.apache.tez.dag.api.TezException; |
| import org.apache.tez.dag.api.UserPayload; |
| import org.apache.tez.dag.api.client.DAGClientServer; |
| import org.apache.tez.dag.app.AppContext; |
| import org.apache.tez.dag.app.ContainerContext; |
| import org.apache.tez.dag.app.ServicePluginLifecycleAbstractService; |
| import org.apache.tez.dag.app.dag.DAG; |
| import org.apache.tez.dag.app.dag.TaskAttempt; |
| import org.apache.tez.dag.app.dag.event.DAGAppMasterEventType; |
| import org.apache.tez.dag.app.dag.event.DAGAppMasterEventUserServiceFatalError; |
| import org.apache.tez.dag.app.dag.event.DAGEventTerminateDag; |
| import org.apache.tez.dag.app.dag.impl.TaskAttemptImpl; |
| import org.apache.tez.dag.app.dag.impl.TaskImpl; |
| import org.apache.tez.dag.app.dag.impl.VertexImpl; |
| import org.apache.tez.dag.app.rm.container.AMContainer; |
| import org.apache.tez.dag.app.rm.container.AMContainerEventAssignTA; |
| import org.apache.tez.dag.app.rm.container.AMContainerEventCompleted; |
| import org.apache.tez.dag.app.rm.container.AMContainerEventType; |
| import org.apache.tez.dag.app.rm.container.AMContainerMap; |
| import org.apache.tez.dag.app.rm.container.AMContainerState; |
| import org.apache.tez.dag.app.web.WebUIService; |
| import org.apache.tez.dag.helpers.DagInfoImplForTest; |
| import org.apache.tez.dag.records.TaskAttemptTerminationCause; |
| import org.apache.tez.dag.records.TezDAGID; |
| import org.apache.tez.dag.records.TezTaskAttemptID; |
| import org.apache.tez.dag.records.TezTaskID; |
| import org.apache.tez.dag.records.TezVertexID; |
| import org.apache.tez.hadoop.shim.HadoopShimsLoader; |
| import org.apache.tez.runtime.api.impl.TaskSpec; |
| import org.apache.tez.serviceplugins.api.ServicePluginErrorDefaults; |
| import org.apache.tez.serviceplugins.api.ServicePluginException; |
| import org.apache.tez.serviceplugins.api.TaskAttemptEndReason; |
| import org.apache.tez.serviceplugins.api.TaskScheduler; |
| import org.apache.tez.serviceplugins.api.TaskSchedulerContext; |
| import org.apache.tez.serviceplugins.api.TaskSchedulerDescriptor; |
| import org.junit.Assert; |
| import org.junit.Before; |
| import org.junit.Test; |
| |
| import com.google.common.collect.Lists; |
| import org.mockito.ArgumentCaptor; |
| import org.mockito.invocation.InvocationOnMock; |
| import org.mockito.stubbing.Answer; |
| |
| @SuppressWarnings("rawtypes") |
| public class TestTaskSchedulerManager { |
| |
| class TestEventHandler implements EventHandler{ |
| List<Event> events = Lists.newLinkedList(); |
| @Override |
| public void handle(Event event) { |
| events.add(event); |
| } |
| } |
| |
| class MockTaskSchedulerManager extends TaskSchedulerManager { |
| |
| final AtomicBoolean notify = new AtomicBoolean(false); |
| |
| public MockTaskSchedulerManager(AppContext appContext, |
| DAGClientServer clientService, EventHandler eventHandler, |
| ContainerSignatureMatcher containerSignatureMatcher, |
| WebUIService webUI) { |
| super(appContext, clientService, eventHandler, containerSignatureMatcher, webUI, |
| Lists.newArrayList(new NamedEntityDescriptor("FakeDescriptor", null)), false, |
| new HadoopShimsLoader(appContext.getAMConf()).getHadoopShim()); |
| } |
| |
| @Override |
| protected void instantiateSchedulers(String host, int port, String trackingUrl, |
| AppContext appContext) { |
| taskSchedulers[0] = new TaskSchedulerWrapper(mockTaskScheduler); |
| taskSchedulerServiceWrappers[0] = |
| new ServicePluginLifecycleAbstractService<>(taskSchedulers[0].getTaskScheduler()); |
| } |
| |
| @Override |
| protected void notifyForTest() { |
| synchronized (notify) { |
| notify.set(true); |
| notify.notifyAll(); |
| } |
| } |
| |
| } |
| |
| AppContext mockAppContext; |
| DAGClientServer mockClientService; |
| TestEventHandler mockEventHandler; |
| ContainerSignatureMatcher mockSigMatcher; |
| MockTaskSchedulerManager schedulerHandler; |
| TaskScheduler mockTaskScheduler; |
| AMContainerMap mockAMContainerMap; |
| WebUIService mockWebUIService; |
| |
| @Before |
| public void setup() { |
| mockAppContext = mock(AppContext.class, RETURNS_DEEP_STUBS); |
| doReturn(new Configuration(false)).when(mockAppContext).getAMConf(); |
| mockClientService = mock(DAGClientServer.class); |
| mockEventHandler = new TestEventHandler(); |
| mockSigMatcher = mock(ContainerSignatureMatcher.class); |
| mockTaskScheduler = mock(TaskScheduler.class); |
| mockAMContainerMap = mock(AMContainerMap.class); |
| mockWebUIService = mock(WebUIService.class); |
| when(mockAppContext.getAllContainers()).thenReturn(mockAMContainerMap); |
| when(mockClientService.getBindAddress()).thenReturn(new InetSocketAddress(10000)); |
| schedulerHandler = new MockTaskSchedulerManager( |
| mockAppContext, mockClientService, mockEventHandler, mockSigMatcher, mockWebUIService); |
| } |
| |
| @Test(timeout = 5000) |
| public void testSimpleAllocate() throws Exception { |
| Configuration conf = new Configuration(false); |
| schedulerHandler.init(conf); |
| schedulerHandler.start(); |
| |
| TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class); |
| TezTaskAttemptID mockAttemptId = mock(TezTaskAttemptID.class); |
| when(mockAttemptId.getId()).thenReturn(0); |
| when(mockTaskAttempt.getID()).thenReturn(mockAttemptId); |
| Resource resource = Resource.newInstance(1024, 1); |
| ContainerContext containerContext = |
| new ContainerContext(new HashMap<String, LocalResource>(), new Credentials(), |
| new HashMap<String, String>(), ""); |
| int priority = 10; |
| TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(new HashSet<String>(), null); |
| |
| ContainerId mockCId = mock(ContainerId.class); |
| Container container = mock(Container.class); |
| when(container.getId()).thenReturn(mockCId); |
| |
| AMContainer mockAMContainer = mock(AMContainer.class); |
| when(mockAMContainer.getContainerId()).thenReturn(mockCId); |
| when(mockAMContainer.getState()).thenReturn(AMContainerState.IDLE); |
| |
| when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); |
| |
| AMSchedulerEventTALaunchRequest lr = |
| new AMSchedulerEventTALaunchRequest(mockAttemptId, resource, null, mockTaskAttempt, locHint, |
| priority, containerContext, 0, 0, 0); |
| schedulerHandler.taskAllocated(0, mockTaskAttempt, lr, container); |
| assertEquals(1, mockEventHandler.events.size()); |
| assertTrue(mockEventHandler.events.get(0) instanceof AMContainerEventAssignTA); |
| AMContainerEventAssignTA assignEvent = |
| (AMContainerEventAssignTA) mockEventHandler.events.get(0); |
| assertEquals(priority, assignEvent.getPriority()); |
| assertEquals(mockAttemptId, assignEvent.getTaskAttemptId()); |
| } |
| |
| @Test (timeout = 5000) |
| public void testTaskBasedAffinity() throws Exception { |
| Configuration conf = new Configuration(false); |
| schedulerHandler.init(conf); |
| schedulerHandler.start(); |
| |
| TaskAttemptImpl mockTaskAttempt = mock(TaskAttemptImpl.class); |
| TezTaskAttemptID taId = mock(TezTaskAttemptID.class); |
| String affVertexName = "srcVertex"; |
| int affTaskIndex = 1; |
| TaskLocationHint locHint = TaskLocationHint.createTaskLocationHint(affVertexName, affTaskIndex); |
| VertexImpl affVertex = mock(VertexImpl.class); |
| TaskImpl affTask = mock(TaskImpl.class); |
| TaskAttemptImpl affAttempt = mock(TaskAttemptImpl.class); |
| ContainerId affCId = mock(ContainerId.class); |
| when(affVertex.getTotalTasks()).thenReturn(2); |
| when(affVertex.getTask(affTaskIndex)).thenReturn(affTask); |
| when(affTask.getSuccessfulAttempt()).thenReturn(affAttempt); |
| when(affAttempt.getAssignedContainerID()).thenReturn(affCId); |
| when(mockAppContext.getCurrentDAG().getVertex(affVertexName)).thenReturn(affVertex); |
| Resource resource = Resource.newInstance(100, 1); |
| AMSchedulerEventTALaunchRequest event = new AMSchedulerEventTALaunchRequest |
| (taId, resource, null, mockTaskAttempt, locHint, 3, null, 0, 0, 0); |
| schedulerHandler.notify.set(false); |
| schedulerHandler.handle(event); |
| synchronized (schedulerHandler.notify) { |
| while (!schedulerHandler.notify.get()) { |
| schedulerHandler.notify.wait(); |
| } |
| } |
| |
| // verify mockTaskAttempt affinitized to expected affCId |
| verify(mockTaskScheduler, times(1)).allocateTask(mockTaskAttempt, resource, affCId, |
| Priority.newInstance(3), null, event); |
| |
| schedulerHandler.stop(); |
| schedulerHandler.close(); |
| } |
| |
| @Test (timeout = 5000) |
| public void testContainerPreempted() throws IOException { |
| Configuration conf = new Configuration(false); |
| schedulerHandler.init(conf); |
| schedulerHandler.start(); |
| |
| String diagnostics = "Container preempted by RM."; |
| TaskAttemptImpl mockTask = mock(TaskAttemptImpl.class); |
| ContainerStatus mockStatus = mock(ContainerStatus.class); |
| ContainerId mockCId = mock(ContainerId.class); |
| AMContainer mockAMContainer = mock(AMContainer.class); |
| when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); |
| when(mockAMContainer.getContainerId()).thenReturn(mockCId); |
| when(mockStatus.getContainerId()).thenReturn(mockCId); |
| when(mockStatus.getDiagnostics()).thenReturn(diagnostics); |
| when(mockStatus.getExitStatus()).thenReturn(ContainerExitStatus.PREEMPTED); |
| schedulerHandler.containerCompleted(0, mockTask, mockStatus); |
| assertEquals(1, mockEventHandler.events.size()); |
| Event event = mockEventHandler.events.get(0); |
| assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); |
| AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; |
| assertEquals(mockCId, completedEvent.getContainerId()); |
| assertEquals("Container preempted externally. Container preempted by RM.", |
| completedEvent.getDiagnostics()); |
| assertTrue(completedEvent.isPreempted()); |
| assertEquals(TaskAttemptTerminationCause.EXTERNAL_PREEMPTION, |
| completedEvent.getTerminationCause()); |
| Assert.assertFalse(completedEvent.isDiskFailed()); |
| |
| schedulerHandler.stop(); |
| schedulerHandler.close(); |
| } |
| |
| @Test (timeout = 5000) |
| public void testContainerInternalPreempted() throws IOException, ServicePluginException { |
| Configuration conf = new Configuration(false); |
| schedulerHandler.init(conf); |
| schedulerHandler.start(); |
| |
| AMContainer mockAmContainer = mock(AMContainer.class); |
| when(mockAmContainer.getTaskSchedulerIdentifier()).thenReturn(0); |
| when(mockAmContainer.getContainerLauncherIdentifier()).thenReturn(0); |
| when(mockAmContainer.getTaskCommunicatorIdentifier()).thenReturn(0); |
| ContainerId mockCId = mock(ContainerId.class); |
| verify(mockTaskScheduler, times(0)).deallocateContainer((ContainerId) any()); |
| when(mockAMContainerMap.get(mockCId)).thenReturn(mockAmContainer); |
| schedulerHandler.preemptContainer(0, mockCId); |
| verify(mockTaskScheduler, times(1)).deallocateContainer(mockCId); |
| assertEquals(1, mockEventHandler.events.size()); |
| Event event = mockEventHandler.events.get(0); |
| assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); |
| AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; |
| assertEquals(mockCId, completedEvent.getContainerId()); |
| assertEquals("Container preempted internally", completedEvent.getDiagnostics()); |
| assertTrue(completedEvent.isPreempted()); |
| Assert.assertFalse(completedEvent.isDiskFailed()); |
| assertEquals(TaskAttemptTerminationCause.INTERNAL_PREEMPTION, |
| completedEvent.getTerminationCause()); |
| |
| schedulerHandler.stop(); |
| schedulerHandler.close(); |
| } |
| |
| @Test (timeout = 5000) |
| public void testContainerDiskFailed() throws IOException { |
| Configuration conf = new Configuration(false); |
| schedulerHandler.init(conf); |
| schedulerHandler.start(); |
| |
| String diagnostics = "NM disk failed."; |
| TaskAttemptImpl mockTask = mock(TaskAttemptImpl.class); |
| ContainerStatus mockStatus = mock(ContainerStatus.class); |
| ContainerId mockCId = mock(ContainerId.class); |
| AMContainer mockAMContainer = mock(AMContainer.class); |
| when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); |
| when(mockAMContainer.getContainerId()).thenReturn(mockCId); |
| when(mockStatus.getContainerId()).thenReturn(mockCId); |
| when(mockStatus.getDiagnostics()).thenReturn(diagnostics); |
| when(mockStatus.getExitStatus()).thenReturn(ContainerExitStatus.DISKS_FAILED); |
| schedulerHandler.containerCompleted(0, mockTask, mockStatus); |
| assertEquals(1, mockEventHandler.events.size()); |
| Event event = mockEventHandler.events.get(0); |
| assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); |
| AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; |
| assertEquals(mockCId, completedEvent.getContainerId()); |
| assertEquals("Container disk failed. NM disk failed.", |
| completedEvent.getDiagnostics()); |
| Assert.assertFalse(completedEvent.isPreempted()); |
| assertTrue(completedEvent.isDiskFailed()); |
| assertEquals(TaskAttemptTerminationCause.NODE_DISK_ERROR, |
| completedEvent.getTerminationCause()); |
| |
| schedulerHandler.stop(); |
| schedulerHandler.close(); |
| } |
| |
| @Test (timeout = 5000) |
| public void testContainerExceededPMem() throws IOException { |
| Configuration conf = new Configuration(false); |
| schedulerHandler.init(conf); |
| schedulerHandler.start(); |
| |
| String diagnostics = "Exceeded Physical Memory"; |
| TaskAttemptImpl mockTask = mock(TaskAttemptImpl.class); |
| ContainerStatus mockStatus = mock(ContainerStatus.class); |
| ContainerId mockCId = mock(ContainerId.class); |
| AMContainer mockAMContainer = mock(AMContainer.class); |
| when(mockAMContainerMap.get(mockCId)).thenReturn(mockAMContainer); |
| when(mockAMContainer.getContainerId()).thenReturn(mockCId); |
| when(mockStatus.getContainerId()).thenReturn(mockCId); |
| when(mockStatus.getDiagnostics()).thenReturn(diagnostics); |
| // use -104 rather than ContainerExitStatus.KILLED_EXCEEDED_PMEM because |
| // ContainerExitStatus.KILLED_EXCEEDED_PMEM is only available after hadoop-2.5 |
| when(mockStatus.getExitStatus()).thenReturn(-104); |
| schedulerHandler.containerCompleted(0, mockTask, mockStatus); |
| assertEquals(1, mockEventHandler.events.size()); |
| Event event = mockEventHandler.events.get(0); |
| assertEquals(AMContainerEventType.C_COMPLETED, event.getType()); |
| AMContainerEventCompleted completedEvent = (AMContainerEventCompleted) event; |
| assertEquals(mockCId, completedEvent.getContainerId()); |
| assertEquals("Container failed, exitCode=-104. Exceeded Physical Memory", |
| completedEvent.getDiagnostics()); |
| Assert.assertFalse(completedEvent.isPreempted()); |
| Assert.assertFalse(completedEvent.isDiskFailed()); |
| assertEquals(TaskAttemptTerminationCause.CONTAINER_EXITED, |
| completedEvent.getTerminationCause()); |
| |
| schedulerHandler.stop(); |
| schedulerHandler.close(); |
| } |
| |
| @Test (timeout = 5000) |
| public void testHistoryUrlConf() throws Exception { |
| Configuration conf = schedulerHandler.appContext.getAMConf(); |
| final ApplicationId mockApplicationId = mock(ApplicationId.class); |
| doReturn("TEST_APP_ID").when(mockApplicationId).toString(); |
| doReturn(mockApplicationId).when(mockAppContext).getApplicationID(); |
| |
| // ensure history url is empty when timeline server is not the logging class |
| conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "http://ui-host:9999"); |
| assertTrue("http://ui-host:9999/#/tez-app/TEST_APP_ID" |
| .equals(schedulerHandler.getHistoryUrl())); |
| |
| // ensure the trailing / in history url is handled |
| conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "http://ui-host:9998/"); |
| assertTrue("http://ui-host:9998/#/tez-app/TEST_APP_ID" |
| .equals(schedulerHandler.getHistoryUrl())); |
| |
| // ensure missing scheme in history url is handled |
| conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "ui-host:9998/"); |
| Assert.assertTrue("http://ui-host:9998/#/tez-app/TEST_APP_ID" |
| .equals(schedulerHandler.getHistoryUrl())); |
| |
| // handle bad template ex without begining / |
| conf.set(TezConfiguration.TEZ_AM_TEZ_UI_HISTORY_URL_TEMPLATE, |
| "__HISTORY_URL_BASE__#/somepath"); |
| assertTrue("http://ui-host:9998/#/somepath" |
| .equals(schedulerHandler.getHistoryUrl())); |
| |
| conf.set(TezConfiguration.TEZ_AM_TEZ_UI_HISTORY_URL_TEMPLATE, |
| "__HISTORY_URL_BASE__?viewPath=tez-app/__APPLICATION_ID__"); |
| conf.set(TezConfiguration.TEZ_HISTORY_URL_BASE, "http://localhost/ui/tez"); |
| assertTrue("http://localhost/ui/tez?viewPath=tez-app/TEST_APP_ID" |
| .equals(schedulerHandler.getHistoryUrl())); |
| |
| } |
| |
| @Test(timeout = 5000) |
| public void testNoSchedulerSpecified() throws IOException { |
| try { |
| new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler, |
| mockSigMatcher, mockWebUIService, null, false); |
| fail("Expecting an IllegalStateException with no schedulers specified"); |
| } catch (IllegalArgumentException e) { |
| } |
| } |
| |
| // Verified via statics |
| @Test(timeout = 5000) |
| public void testCustomTaskSchedulerSetup() throws IOException { |
| Configuration conf = new Configuration(false); |
| conf.set("testkey", "testval"); |
| UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf); |
| |
| String customSchedulerName = "fakeScheduler"; |
| List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>(); |
| ByteBuffer bb = ByteBuffer.allocate(4); |
| bb.putInt(0, 3); |
| UserPayload userPayload = UserPayload.create(bb); |
| taskSchedulers.add( |
| new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName()) |
| .setUserPayload(userPayload)); |
| taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null) |
| .setUserPayload(defaultPayload)); |
| |
| TSEHForMultipleSchedulersTest tseh = |
| new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler, |
| mockSigMatcher, mockWebUIService, taskSchedulers, false); |
| |
| tseh.init(conf); |
| tseh.start(); |
| |
| // Verify that the YARN task scheduler is installed by default |
| assertTrue(tseh.getYarnSchedulerCreated()); |
| assertFalse(tseh.getUberSchedulerCreated()); |
| assertEquals(2, tseh.getNumCreateInvocations()); |
| |
| // Verify the order of the schedulers |
| assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0)); |
| assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1)); |
| |
| // Verify the payload setup for the custom task scheduler |
| assertNotNull(tseh.getTaskSchedulerContext(0)); |
| assertEquals(bb, tseh.getTaskSchedulerContext(0).getInitialUserPayload().getPayload()); |
| |
| // Verify the payload on the yarn scheduler |
| assertNotNull(tseh.getTaskSchedulerContext(1)); |
| Configuration parsed = TezUtils.createConfFromUserPayload(tseh.getTaskSchedulerContext(1).getInitialUserPayload()); |
| assertEquals("testval", parsed.get("testkey")); |
| } |
| |
| @Test(timeout = 5000) |
| public void testTaskSchedulerRouting() throws Exception { |
| Configuration conf = new Configuration(false); |
| UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(conf); |
| |
| String customSchedulerName = "fakeScheduler"; |
| List<NamedEntityDescriptor> taskSchedulers = new LinkedList<>(); |
| ByteBuffer bb = ByteBuffer.allocate(4); |
| bb.putInt(0, 3); |
| UserPayload userPayload = UserPayload.create(bb); |
| taskSchedulers.add( |
| new NamedEntityDescriptor(customSchedulerName, FakeTaskScheduler.class.getName()) |
| .setUserPayload(userPayload)); |
| taskSchedulers.add(new NamedEntityDescriptor(TezConstants.getTezYarnServicePluginName(), null) |
| .setUserPayload(defaultPayload)); |
| |
| TSEHForMultipleSchedulersTest tseh = |
| new TSEHForMultipleSchedulersTest(mockAppContext, mockClientService, mockEventHandler, |
| mockSigMatcher, mockWebUIService, taskSchedulers, false); |
| |
| tseh.init(conf); |
| tseh.start(); |
| |
| // Verify that the YARN task scheduler is installed by default |
| assertTrue(tseh.getYarnSchedulerCreated()); |
| assertFalse(tseh.getUberSchedulerCreated()); |
| assertEquals(2, tseh.getNumCreateInvocations()); |
| |
| // Verify the order of the schedulers |
| assertEquals(customSchedulerName, tseh.getTaskSchedulerName(0)); |
| assertEquals(TezConstants.getTezYarnServicePluginName(), tseh.getTaskSchedulerName(1)); |
| |
| verify(tseh.getTestTaskScheduler(0)).initialize(); |
| verify(tseh.getTestTaskScheduler(0)).start(); |
| |
| ApplicationId appId = ApplicationId.newInstance(1000, 1); |
| TezDAGID dagId = TezDAGID.getInstance(appId, 1); |
| TezVertexID vertexID = TezVertexID.getInstance(dagId, 1); |
| TezTaskID taskId1 = TezTaskID.getInstance(vertexID, 1); |
| TezTaskAttemptID attemptId11 = TezTaskAttemptID.getInstance(taskId1, 1); |
| TezTaskID taskId2 = TezTaskID.getInstance(vertexID, 2); |
| TezTaskAttemptID attemptId21 = TezTaskAttemptID.getInstance(taskId2, 1); |
| |
| Resource resource = Resource.newInstance(1024, 1); |
| |
| TaskAttempt mockTaskAttempt1 = mock(TaskAttempt.class); |
| TaskAttempt mockTaskAttempt2 = mock(TaskAttempt.class); |
| |
| AMSchedulerEventTALaunchRequest launchRequest1 = |
| new AMSchedulerEventTALaunchRequest(attemptId11, resource, mock(TaskSpec.class), |
| mockTaskAttempt1, mock(TaskLocationHint.class), 1, mock(ContainerContext.class), 0, 0, |
| 0); |
| |
| tseh.handle(launchRequest1); |
| |
| verify(tseh.getTestTaskScheduler(0)).allocateTask(eq(mockTaskAttempt1), eq(resource), |
| any(String[].class), any(String[].class), any(Priority.class), any(Object.class), |
| eq(launchRequest1)); |
| |
| AMSchedulerEventTALaunchRequest launchRequest2 = |
| new AMSchedulerEventTALaunchRequest(attemptId21, resource, mock(TaskSpec.class), |
| mockTaskAttempt2, mock(TaskLocationHint.class), 1, mock(ContainerContext.class), 1, 0, |
| 0); |
| tseh.handle(launchRequest2); |
| verify(tseh.getTestTaskScheduler(1)).allocateTask(eq(mockTaskAttempt2), eq(resource), |
| any(String[].class), any(String[].class), any(Priority.class), any(Object.class), |
| eq(launchRequest2)); |
| } |
| |
| @SuppressWarnings("unchecked") |
| @Test(timeout = 5000) |
| public void testReportFailureFromTaskScheduler() { |
| String dagName = DAG_NAME; |
| Configuration conf = new TezConfiguration(); |
| String taskSchedulerName = "testTaskScheduler"; |
| String expIdentifier = "[0:" + taskSchedulerName + "]"; |
| EventHandler eventHandler = mock(EventHandler.class); |
| AppContext appContext = mock(AppContext.class, RETURNS_DEEP_STUBS); |
| doReturn(taskSchedulerName).when(appContext).getTaskSchedulerName(0); |
| doReturn(eventHandler).when(appContext).getEventHandler(); |
| doReturn(conf).when(appContext).getAMConf(); |
| InetSocketAddress address = new InetSocketAddress("host", 55000); |
| |
| DAGClientServer dagClientServer = mock(DAGClientServer.class); |
| doReturn(address).when(dagClientServer).getBindAddress(); |
| |
| DAG dag = mock(DAG.class); |
| TezDAGID dagId = TezDAGID.getInstance(ApplicationId.newInstance(1, 0), DAG_INDEX); |
| doReturn(dagName).when(dag).getName(); |
| doReturn(dagId).when(dag).getID(); |
| doReturn(dag).when(appContext).getCurrentDAG(); |
| |
| NamedEntityDescriptor<TaskSchedulerDescriptor> namedEntityDescriptor = |
| new NamedEntityDescriptor<>(taskSchedulerName, TaskSchedulerForFailureTest.class.getName()); |
| List<NamedEntityDescriptor> list = new LinkedList<>(); |
| list.add(namedEntityDescriptor); |
| |
| TaskSchedulerManager taskSchedulerManager = |
| new TaskSchedulerManager(appContext, dagClientServer, eventHandler, |
| mock(ContainerSignatureMatcher.class), mock(WebUIService.class), list, false, |
| new HadoopShimsLoader(appContext.getAMConf()).getHadoopShim()) { |
| @Override |
| TaskSchedulerContext wrapTaskSchedulerContext(TaskSchedulerContext rawContext) { |
| // Avoid wrapping in threads |
| return rawContext; |
| } |
| }; |
| try { |
| taskSchedulerManager.init(new TezConfiguration()); |
| taskSchedulerManager.start(); |
| |
| taskSchedulerManager.getTotalResources(0); |
| ArgumentCaptor<Event> argumentCaptor = ArgumentCaptor.forClass(Event.class); |
| verify(eventHandler, times(1)).handle(argumentCaptor.capture()); |
| |
| Event rawEvent = argumentCaptor.getValue(); |
| assertTrue(rawEvent instanceof DAGEventTerminateDag); |
| DAGEventTerminateDag killEvent = (DAGEventTerminateDag) rawEvent; |
| assertTrue(killEvent.getDiagnosticInfo().contains("ReportError")); |
| assertTrue(killEvent.getDiagnosticInfo() |
| .contains(ServicePluginErrorDefaults.SERVICE_UNAVAILABLE.name())); |
| assertTrue(killEvent.getDiagnosticInfo().contains(expIdentifier)); |
| |
| |
| reset(eventHandler); |
| taskSchedulerManager.getAvailableResources(0); |
| argumentCaptor = ArgumentCaptor.forClass(Event.class); |
| |
| verify(eventHandler, times(1)).handle(argumentCaptor.capture()); |
| rawEvent = argumentCaptor.getValue(); |
| |
| assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError); |
| DAGAppMasterEventUserServiceFatalError event = |
| (DAGAppMasterEventUserServiceFatalError) rawEvent; |
| assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType()); |
| assertTrue(event.getDiagnosticInfo().contains("ReportedFatalError")); |
| assertTrue( |
| event.getDiagnosticInfo().contains(ServicePluginErrorDefaults.INCONSISTENT_STATE.name())); |
| assertTrue(event.getDiagnosticInfo().contains(expIdentifier)); |
| |
| } finally { |
| taskSchedulerManager.stop(); |
| } |
| } |
| |
| @SuppressWarnings("unchecked") |
| @Test(timeout = 5000) |
| public void testTaskSchedulerUserError() { |
| TaskScheduler taskScheduler = mock(TaskScheduler.class, new ExceptionAnswer()); |
| |
| EventHandler eventHandler = mock(EventHandler.class); |
| AppContext appContext = mock(AppContext.class, RETURNS_DEEP_STUBS); |
| |
| when(appContext.getEventHandler()).thenReturn(eventHandler); |
| doReturn("testTaskScheduler").when(appContext).getTaskSchedulerName(0); |
| String expectedId = "[0:testTaskScheduler]"; |
| |
| Configuration conf = new Configuration(false); |
| |
| InetSocketAddress address = new InetSocketAddress(15222); |
| DAGClientServer mockClientService = mock(DAGClientServer.class); |
| doReturn(address).when(mockClientService).getBindAddress(); |
| TaskSchedulerManager taskSchedulerManager = |
| new TaskSchedulerManager(taskScheduler, appContext, mock(ContainerSignatureMatcher.class), |
| mockClientService, Executors.newFixedThreadPool(1)) { |
| @Override |
| protected void instantiateSchedulers(String host, int port, String trackingUrl, |
| AppContext appContext) throws TezException { |
| // Stubbed out since these are setup up front in the constructor used for testing |
| } |
| }; |
| |
| try { |
| taskSchedulerManager.init(conf); |
| taskSchedulerManager.start(); |
| |
| // Invoking a couple of random methods |
| |
| AMSchedulerEventTALaunchRequest launchRequest = |
| new AMSchedulerEventTALaunchRequest(mock(TezTaskAttemptID.class), mock(Resource.class), |
| mock(TaskSpec.class), mock(TaskAttempt.class), mock(TaskLocationHint.class), 0, |
| mock(ContainerContext.class), 0, 0, 0); |
| taskSchedulerManager.handleEvent(launchRequest); |
| |
| ArgumentCaptor<Event> argumentCaptor = ArgumentCaptor.forClass(Event.class); |
| |
| verify(eventHandler, times(1)).handle(argumentCaptor.capture()); |
| |
| Event rawEvent = argumentCaptor.getValue(); |
| assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError); |
| DAGAppMasterEventUserServiceFatalError event = |
| (DAGAppMasterEventUserServiceFatalError) rawEvent; |
| |
| assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType()); |
| assertTrue(event.getError().getMessage().contains("TestException_" + "allocateTask")); |
| assertTrue(event.getDiagnosticInfo().contains("Task Allocation")); |
| assertTrue(event.getDiagnosticInfo().contains(expectedId)); |
| |
| |
| taskSchedulerManager.dagCompleted(); |
| argumentCaptor = ArgumentCaptor.forClass(Event.class); |
| verify(eventHandler, times(2)).handle(argumentCaptor.capture()); |
| |
| rawEvent = argumentCaptor.getAllValues().get(1); |
| assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError); |
| event = (DAGAppMasterEventUserServiceFatalError) rawEvent; |
| |
| assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType()); |
| assertTrue(event.getError().getMessage().contains("TestException_" + "dagComplete")); |
| assertTrue(event.getDiagnosticInfo().contains("Dag Completion")); |
| assertTrue(event.getDiagnosticInfo().contains(expectedId)); |
| |
| } finally { |
| taskSchedulerManager.stop(); |
| } |
| } |
| |
| private static class ExceptionAnswer implements Answer { |
| @Override |
| public Object answer(InvocationOnMock invocation) throws Throwable { |
| Method method = invocation.getMethod(); |
| if (method.getDeclaringClass().equals(TaskScheduler.class) && |
| !method.getName().equals("getContext") && !method.getName().equals("initialize") && |
| !method.getName().equals("start") && !method.getName().equals("shutdown")) { |
| throw new RuntimeException("TestException_" + method.getName()); |
| } else { |
| return invocation.callRealMethod(); |
| } |
| } |
| } |
| |
| public static class TSEHForMultipleSchedulersTest extends TaskSchedulerManager { |
| |
| private final TaskScheduler yarnTaskScheduler; |
| private final TaskScheduler uberTaskScheduler; |
| private final AtomicBoolean uberSchedulerCreated = new AtomicBoolean(false); |
| private final AtomicBoolean yarnSchedulerCreated = new AtomicBoolean(false); |
| private final AtomicInteger numCreateInvocations = new AtomicInteger(0); |
| private final Set<Integer> seenSchedulers = new HashSet<>(); |
| private final List<TaskSchedulerContext> taskSchedulerContexts = new LinkedList<>(); |
| private final List<String> taskSchedulerNames = new LinkedList<>(); |
| private final List<TaskScheduler> testTaskSchedulers = new LinkedList<>(); |
| |
| public TSEHForMultipleSchedulersTest(AppContext appContext, |
| DAGClientServer clientService, |
| EventHandler eventHandler, |
| ContainerSignatureMatcher containerSignatureMatcher, |
| WebUIService webUI, |
| List<NamedEntityDescriptor> schedulerDescriptors, |
| boolean isPureLocalMode) { |
| super(appContext, clientService, eventHandler, containerSignatureMatcher, webUI, |
| schedulerDescriptors, isPureLocalMode, |
| new HadoopShimsLoader(appContext.getAMConf()).getHadoopShim()); |
| yarnTaskScheduler = mock(TaskScheduler.class); |
| uberTaskScheduler = mock(TaskScheduler.class); |
| } |
| |
| @Override |
| TaskScheduler createTaskScheduler(String host, int port, String trackingUrl, |
| AppContext appContext, |
| NamedEntityDescriptor taskSchedulerDescriptor, |
| long customAppIdIdentifier, |
| int schedulerId) throws TezException { |
| |
| numCreateInvocations.incrementAndGet(); |
| boolean added = seenSchedulers.add(schedulerId); |
| assertTrue("Cannot add multiple schedulers with the same schedulerId", added); |
| taskSchedulerNames.add(taskSchedulerDescriptor.getEntityName()); |
| return super.createTaskScheduler(host, port, trackingUrl, appContext, taskSchedulerDescriptor, |
| customAppIdIdentifier, schedulerId); |
| } |
| |
| @Override |
| TaskSchedulerContext wrapTaskSchedulerContext(TaskSchedulerContext rawContext) { |
| // Avoid wrapping in threads |
| return rawContext; |
| } |
| |
| @Override |
| TaskScheduler createYarnTaskScheduler(TaskSchedulerContext taskSchedulerContext, int schedulerId) { |
| taskSchedulerContexts.add(taskSchedulerContext); |
| testTaskSchedulers.add(yarnTaskScheduler); |
| yarnSchedulerCreated.set(true); |
| return yarnTaskScheduler; |
| } |
| |
| @Override |
| TaskScheduler createUberTaskScheduler(TaskSchedulerContext taskSchedulerContext, int schedulerId) { |
| taskSchedulerContexts.add(taskSchedulerContext); |
| uberSchedulerCreated.set(true); |
| testTaskSchedulers.add(yarnTaskScheduler); |
| return uberTaskScheduler; |
| } |
| |
| @Override |
| TaskScheduler createCustomTaskScheduler(TaskSchedulerContext taskSchedulerContext, |
| NamedEntityDescriptor taskSchedulerDescriptor, int schedulerId) |
| throws TezException { |
| taskSchedulerContexts.add(taskSchedulerContext); |
| TaskScheduler taskScheduler = spy(super.createCustomTaskScheduler(taskSchedulerContext, taskSchedulerDescriptor, schedulerId)); |
| testTaskSchedulers.add(taskScheduler); |
| return taskScheduler; |
| } |
| |
| @Override |
| // Inline handling of events. |
| public void handle(AMSchedulerEvent event) { |
| handleEvent(event); |
| } |
| |
| public boolean getUberSchedulerCreated() { |
| return uberSchedulerCreated.get(); |
| } |
| |
| public boolean getYarnSchedulerCreated() { |
| return yarnSchedulerCreated.get(); |
| } |
| |
| public int getNumCreateInvocations() { |
| return numCreateInvocations.get(); |
| } |
| |
| public TaskSchedulerContext getTaskSchedulerContext(int schedulerId) { |
| return taskSchedulerContexts.get(schedulerId); |
| } |
| |
| public String getTaskSchedulerName(int schedulerId) { |
| return taskSchedulerNames.get(schedulerId); |
| } |
| |
| public TaskScheduler getTestTaskScheduler(int schedulerId) { |
| return testTaskSchedulers.get(schedulerId); |
| } |
| } |
| |
| public static class FakeTaskScheduler extends TaskScheduler { |
| |
| public FakeTaskScheduler( |
| TaskSchedulerContext taskSchedulerContext) { |
| super(taskSchedulerContext); |
| } |
| |
| @Override |
| public Resource getAvailableResources() { |
| return null; |
| } |
| |
| @Override |
| public int getClusterNodeCount() { |
| return 0; |
| } |
| |
| @Override |
| public void dagComplete() { |
| |
| } |
| |
| @Override |
| public Resource getTotalResources() { |
| return null; |
| } |
| |
| @Override |
| public void blacklistNode(NodeId nodeId) { |
| |
| } |
| |
| @Override |
| public void unblacklistNode(NodeId nodeId) { |
| |
| } |
| |
| @Override |
| public void allocateTask(Object task, Resource capability, String[] hosts, String[] racks, |
| Priority priority, Object containerSignature, Object clientCookie) { |
| |
| } |
| |
| @Override |
| public void allocateTask(Object task, Resource capability, ContainerId containerId, |
| Priority priority, Object containerSignature, Object clientCookie) { |
| |
| } |
| |
| @Override |
| public boolean deallocateTask(Object task, boolean taskSucceeded, |
| TaskAttemptEndReason endReason, |
| String diagnostics) { |
| return false; |
| } |
| |
| @Override |
| public Object deallocateContainer(ContainerId containerId) { |
| return null; |
| } |
| |
| @Override |
| public void setShouldUnregister() { |
| |
| } |
| |
| @Override |
| public boolean hasUnregistered() { |
| return false; |
| } |
| } |
| |
| private static final String DAG_NAME = "dagName"; |
| private static final int DAG_INDEX = 1; |
| public static class TaskSchedulerForFailureTest extends TaskScheduler { |
| |
| public TaskSchedulerForFailureTest(TaskSchedulerContext taskSchedulerContext) { |
| super(taskSchedulerContext); |
| } |
| |
| @Override |
| public Resource getAvailableResources() throws ServicePluginException { |
| getContext().reportError(ServicePluginErrorDefaults.INCONSISTENT_STATE, "ReportedFatalError", null); |
| return Resource.newInstance(1024, 1); |
| } |
| |
| @Override |
| public Resource getTotalResources() throws ServicePluginException { |
| getContext() |
| .reportError(ServicePluginErrorDefaults.SERVICE_UNAVAILABLE, "ReportError", new DagInfoImplForTest(DAG_INDEX, DAG_NAME)); |
| return Resource.newInstance(1024, 1); |
| } |
| |
| @Override |
| public int getClusterNodeCount() throws ServicePluginException { |
| return 0; |
| } |
| |
| @Override |
| public void blacklistNode(NodeId nodeId) throws ServicePluginException { |
| |
| } |
| |
| @Override |
| public void unblacklistNode(NodeId nodeId) throws ServicePluginException { |
| |
| } |
| |
| @Override |
| public void allocateTask(Object task, Resource capability, String[] hosts, String[] racks, |
| Priority priority, Object containerSignature, |
| Object clientCookie) throws |
| ServicePluginException { |
| |
| } |
| |
| @Override |
| public void allocateTask(Object task, Resource capability, ContainerId containerId, |
| Priority priority, Object containerSignature, |
| Object clientCookie) throws |
| ServicePluginException { |
| |
| } |
| |
| @Override |
| public boolean deallocateTask(Object task, boolean taskSucceeded, |
| TaskAttemptEndReason endReason, |
| @Nullable String diagnostics) throws ServicePluginException { |
| return false; |
| } |
| |
| @Override |
| public Object deallocateContainer(ContainerId containerId) throws ServicePluginException { |
| return null; |
| } |
| |
| @Override |
| public void setShouldUnregister() throws ServicePluginException { |
| |
| } |
| |
| @Override |
| public boolean hasUnregistered() throws ServicePluginException { |
| return false; |
| } |
| |
| @Override |
| public void dagComplete() throws ServicePluginException { |
| } |
| } |
| } |
