@SuppressWarnings("unchecked")
@Override
public void handle(VertexEvent event) {
DAG dag = context.getCurrentDAG();
org.apache.tez.dag.app.dag.Vertex vertex = dag.getVertex(event.getVertexId());
((EventHandler<VertexEvent>) vertex).handle(event);
}
@SuppressWarnings("unchecked")
@Override
public void handle(TaskAttemptEvent event) {
DAG dag = context.getCurrentDAG();
Task task =
dag.getVertex(event.getTaskAttemptID().getTaskID().getVertexID())
.getTask(event.getTaskAttemptID().getTaskID());
TaskAttempt attempt = task.getAttempt(event.getTaskAttemptID());
((EventHandler<TaskAttemptEvent>) attempt).handle(event);
}
synchronized void shutdownTezAM() {
sessionStopped.set(true);
if (currentDAG != null && !currentDAG.isComplete()) {
// send a DAG_KILL message
LOG.info("Sending a kill event to the current DAG" + ", dagId=" + currentDAG.getID());
sendEvent(new DAGEvent(currentDAG.getID(), DAGEventType.DAG_KILL));
} else {
LOG.info("No current running DAG, shutting down the AM");
if (isSession && !state.equals(DAGAppMasterState.ERROR)) {
state = DAGAppMasterState.SUCCEEDED;
}
shutdownHandler.shutdown();
}
}
synchronized String submitDAGToAppMaster(DAGPlan dagPlan) throws TezException {
if (currentDAG != null && !state.equals(DAGAppMasterState.IDLE)) {
throw new TezException("App master already running a DAG");
}
if (state.equals(DAGAppMasterState.ERROR) || sessionStopped.get()) {
throw new TezException(
"AM unable to accept new DAG submissions." + " In the process of shutting down");
}
// RPC server runs in the context of the job user as it was started in
// the job user's UGI context
LOG.info("Starting DAG submitted via RPC");
if (LOG.isDebugEnabled()) {
LOG.debug("Writing DAG plan to: " + TezConfiguration.TEZ_PB_PLAN_TEXT_NAME);
File outFile = new File(TezConfiguration.TEZ_PB_PLAN_TEXT_NAME);
try {
PrintWriter printWriter = new PrintWriter(outFile);
String dagPbString = dagPlan.toString();
printWriter.println(dagPbString);
printWriter.close();
} catch (IOException e) {
throw new TezException("Failed to write TEZ_PLAN to " + outFile.toString(), e);
}
}
submittedDAGs.incrementAndGet();
startDAG(dagPlan);
return currentDAG.getID().toString();
}
private void startDAG(DAG dag) {
currentDAG = dag;
this.state = DAGAppMasterState.RUNNING;
// End of creating the job.
((RunningAppContext) context).setDAG(currentDAG);
// create a job event for job initialization
DAGEvent initDagEvent = new DAGEvent(currentDAG.getID(), DAGEventType.DAG_INIT);
// Send init to the job (this does NOT trigger job execution)
// This is a synchronous call, not an event through dispatcher. We want
// job-init to be done completely here.
dagEventDispatcher.handle(initDagEvent);
// All components have started, start the job.
/** create a job-start event to get this ball rolling */
DAGEvent startDagEvent = new DAGEvent(currentDAG.getID(), DAGEventType.DAG_START);
/** send the job-start event. this triggers the job execution. */
sendEvent(startDagEvent);
}
@Override
public TezDAGID getCurrentDAGID() {
try {
rLock.lock();
if (dag != null) {
return dag.getID();
}
return null;
} finally {
rLock.unlock();
}
}
@Test(timeout = 10000)
public void testDAGSchedulerNaturalOrder() {
DAG mockDag = mock(DAG.class);
Vertex mockVertex = mock(Vertex.class);
TaskAttempt mockAttempt = mock(TaskAttempt.class);
when(mockDag.getVertex((TezVertexID) any())).thenReturn(mockVertex);
when(mockVertex.getDistanceFromRoot()).thenReturn(0).thenReturn(1).thenReturn(2);
when(mockAttempt.getIsRescheduled()).thenReturn(false);
DAGEventSchedulerUpdate event =
new DAGEventSchedulerUpdate(DAGEventSchedulerUpdate.UpdateType.TA_SCHEDULE, mockAttempt);
DAGScheduler scheduler = new DAGSchedulerNaturalOrder(mockDag, mockEventHandler);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 4);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
when(mockAttempt.getIsRescheduled()).thenReturn(true);
scheduler.scheduleTask(event);
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
}
public List<String> getDiagnostics() {
if (!isSession) {
if (currentDAG != null) {
return currentDAG.getDiagnostics();
}
} else {
return Collections.singletonList(
"Session stats:"
+ "submittedDAGs="
+ submittedDAGs.get()
+ ", successfulDAGs="
+ successfulDAGs.get()
+ ", failedDAGs="
+ failedDAGs.get());
}
return null;
}
private synchronized void setStateOnDAGCompletion() {
DAGAppMasterState oldState = state;
if (isSession) {
return;
}
switch (currentDAG.getState()) {
case SUCCEEDED:
state = DAGAppMasterState.SUCCEEDED;
break;
case FAILED:
state = DAGAppMasterState.FAILED;
break;
case KILLED:
state = DAGAppMasterState.KILLED;
break;
case ERROR:
state = DAGAppMasterState.ERROR;
break;
default:
state = DAGAppMasterState.ERROR;
break;
}
LOG.info("On DAG completion. Old state: " + oldState + " new state: " + state);
}
@Ignore
@Test(timeout = 10000)
public void testDAGSchedulerMRR() {
DAG mockDag = mock(DAG.class);
TezDAGID dagId = TezDAGID.getInstance("1", 1, 1);
TaskSchedulerEventHandler mockTaskScheduler = mock(TaskSchedulerEventHandler.class);
Vertex mockVertex1 = mock(Vertex.class);
TezVertexID mockVertexId1 = TezVertexID.getInstance(dagId, 1);
when(mockVertex1.getVertexId()).thenReturn(mockVertexId1);
when(mockVertex1.getDistanceFromRoot()).thenReturn(0);
TaskAttempt mockAttempt1 = mock(TaskAttempt.class);
when(mockAttempt1.getVertexID()).thenReturn(mockVertexId1);
when(mockAttempt1.getIsRescheduled()).thenReturn(false);
when(mockDag.getVertex(mockVertexId1)).thenReturn(mockVertex1);
Vertex mockVertex2 = mock(Vertex.class);
TezVertexID mockVertexId2 = TezVertexID.getInstance(dagId, 2);
when(mockVertex2.getVertexId()).thenReturn(mockVertexId2);
when(mockVertex2.getDistanceFromRoot()).thenReturn(1);
TaskAttempt mockAttempt2 = mock(TaskAttempt.class);
when(mockAttempt2.getVertexID()).thenReturn(mockVertexId2);
when(mockAttempt2.getIsRescheduled()).thenReturn(false);
when(mockDag.getVertex(mockVertexId2)).thenReturn(mockVertex2);
TaskAttempt mockAttempt2f = mock(TaskAttempt.class);
when(mockAttempt2f.getVertexID()).thenReturn(mockVertexId2);
when(mockAttempt2f.getIsRescheduled()).thenReturn(true);
Vertex mockVertex3 = mock(Vertex.class);
TezVertexID mockVertexId3 = TezVertexID.getInstance(dagId, 3);
when(mockVertex3.getVertexId()).thenReturn(mockVertexId3);
when(mockVertex3.getDistanceFromRoot()).thenReturn(2);
TaskAttempt mockAttempt3 = mock(TaskAttempt.class);
when(mockAttempt3.getVertexID()).thenReturn(mockVertexId3);
when(mockAttempt3.getIsRescheduled()).thenReturn(false);
when(mockDag.getVertex(mockVertexId3)).thenReturn(mockVertex3);
DAGEventSchedulerUpdate mockEvent1 = mock(DAGEventSchedulerUpdate.class);
when(mockEvent1.getAttempt()).thenReturn(mockAttempt1);
DAGEventSchedulerUpdate mockEvent2 = mock(DAGEventSchedulerUpdate.class);
when(mockEvent2.getAttempt()).thenReturn(mockAttempt2);
DAGEventSchedulerUpdate mockEvent2f = mock(DAGEventSchedulerUpdate.class);
when(mockEvent2f.getAttempt()).thenReturn(mockAttempt2f);
DAGEventSchedulerUpdate mockEvent3 = mock(DAGEventSchedulerUpdate.class);
when(mockEvent3.getAttempt()).thenReturn(mockAttempt3);
DAGScheduler scheduler =
new DAGSchedulerMRR(mockDag, mockEventHandler, mockTaskScheduler, 0.5f);
// M starts. M completes. R1 starts. R1 completes. R2 starts. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent1); // M runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent2); // R1 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent2); // R1 runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent2f); // R1 runs retry. Retry priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 4);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.scheduleTask(mockEvent3); // R2 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 9);
scheduler.scheduleTask(mockEvent3); // R2 runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 9);
scheduler.vertexCompleted(mockVertex3); // R2 completes
// M starts. R1 starts. M completes. R2 starts. R1 completes. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent2); // R1 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.scheduleTask(mockEvent1); // M runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent2); // R1 runs another. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.scheduleTask(mockEvent2f); // R1 runs retry. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent3); // R2 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
scheduler.scheduleTask(mockEvent2); // R1 runs another. Normal priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent2f); // R1 runs retry. Retry priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 4);
scheduler.scheduleTask(mockEvent3); // R2 runs another. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.vertexCompleted(mockVertex3); // R2 completes
// M starts. M completes. R1 starts. R2 starts. R1 completes. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent2); // R1 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.scheduleTask(mockEvent3); // R2 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 5);
scheduler.scheduleTask(mockEvent2); // R1 runs another
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.vertexCompleted(mockVertex3); // R2 completes
// M starts. R1 starts. M completes. R1 completes. R2 starts. R2 completes
scheduler.scheduleTask(mockEvent1); // M starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 3);
scheduler.scheduleTask(mockEvent2); // R1 starts. Reordered priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 2);
scheduler.vertexCompleted(mockVertex1); // M completes
scheduler.scheduleTask(mockEvent2); // R1 starts. Normal priority
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 6);
scheduler.vertexCompleted(mockVertex2); // R1 completes
scheduler.scheduleTask(mockEvent3); // R2 starts
Assert.assertTrue(mockEventHandler.event.getPriority().getPriority() == 9);
scheduler.vertexCompleted(mockVertex3); // R2 completes
}
@Test(timeout = 5000)
public void testDAGClientHandler() throws TezException {
TezDAGID mockTezDAGId = mock(TezDAGID.class);
when(mockTezDAGId.getId()).thenReturn(1);
when(mockTezDAGId.toString()).thenReturn("dag_9999_0001_1");
DAG mockDAG = mock(DAG.class);
when(mockDAG.getID()).thenReturn(mockTezDAGId);
DAGStatusBuilder mockDagStatusBuilder = mock(DAGStatusBuilder.class);
when(mockDAG.getDAGStatus(anySetOf(StatusGetOpts.class))).thenReturn(mockDagStatusBuilder);
VertexStatusBuilder mockVertexStatusBuilder = mock(VertexStatusBuilder.class);
when(mockDAG.getVertexStatus(anyString(), anySetOf(StatusGetOpts.class)))
.thenReturn(mockVertexStatusBuilder);
DAGAppMaster mockDagAM = mock(DAGAppMaster.class);
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.RUNNING);
AppContext mockAppContext = mock(AppContext.class);
when(mockDagAM.getContext()).thenReturn(mockAppContext);
when(mockDagAM.getContext().getCurrentDAG()).thenReturn(mockDAG);
DAGClientHandler dagClientHandler = new DAGClientHandler(mockDagAM);
// getAllDAGs()
assertEquals(1, dagClientHandler.getAllDAGs().size());
assertEquals("dag_9999_0001_1", dagClientHandler.getAllDAGs().get(0));
// getDAGStatus
try {
dagClientHandler.getDAGStatus("dag_9999_0001_2", Sets.newSet(StatusGetOpts.GET_COUNTERS));
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
DAGStatus dagStatus =
dagClientHandler.getDAGStatus("dag_9999_0001_1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
assertEquals(mockDagStatusBuilder, dagStatus);
// getVertexStatus
try {
dagClientHandler.getVertexStatus(
"dag_9999_0001_2", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
VertexStatus vertexStatus =
dagClientHandler.getVertexStatus(
"dag_9999_0001_1", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
assertEquals(mockVertexStatusBuilder, vertexStatus);
// getTezAppMasterStatus
when(mockDagAM.isSession()).thenReturn(false);
assertEquals(TezAppMasterStatus.RUNNING, dagClientHandler.getTezAppMasterStatus());
when(mockDagAM.isSession()).thenReturn(true);
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.INITED);
assertEquals(TezAppMasterStatus.INITIALIZING, dagClientHandler.getTezAppMasterStatus());
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.ERROR);
assertEquals(TezAppMasterStatus.SHUTDOWN, dagClientHandler.getTezAppMasterStatus());
// tryKillDAG
try {
dagClientHandler.tryKillDAG("dag_9999_0001_2");
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
dagClientHandler.tryKillDAG("dag_9999_0001_1");
ArgumentCaptor<DAG> eventCaptor = ArgumentCaptor.forClass(DAG.class);
verify(mockDagAM, times(1))
.tryKillDAG(eventCaptor.capture(), contains("Sending client kill from"));
assertEquals(1, eventCaptor.getAllValues().size());
assertTrue(eventCaptor.getAllValues().get(0) instanceof DAG);
assertEquals("dag_9999_0001_1", ((DAG) eventCaptor.getAllValues().get(0)).getID().toString());
// submitDAG
DAGPlan dagPlan = DAGPlan.getDefaultInstance();
Map<String, LocalResource> localResources = new HashMap<String, LocalResource>();
dagClientHandler.submitDAG(dagPlan, localResources);
verify(mockDagAM).submitDAGToAppMaster(dagPlan, localResources);
// shutdown
dagClientHandler.shutdownAM();
verify(mockDagAM).shutdownTezAM(contains("Received message to shutdown AM from"));
}
@Override
public String getUser() {
return dag.getUserName();
}
public void tryKillDAG(String dagIdStr) throws TezException {
DAG dag = getDAG(dagIdStr);
LOG.info("Sending client kill to dag: " + dagIdStr);
// send a DAG_KILL message
sendEvent(new DAGEvent(dag.getID(), DAGEventType.DAG_KILL));
}
public float getProgress() {
if (currentDAG != null && currentDAG.getState() == DAGState.RUNNING) {
return currentDAG.getProgress();
}
return 0;
}
private synchronized void handle(DAGAppMasterEvent event) {
switch (event.getType()) {
case INTERNAL_ERROR:
state = DAGAppMasterState.ERROR;
if (currentDAG != null) {
// notify dag to finish which will send the DAG_FINISHED event
LOG.info("Internal Error. Notifying dags to finish.");
sendEvent(new DAGEvent(currentDAG.getID(), DAGEventType.INTERNAL_ERROR));
} else {
LOG.info("Internal Error. Finishing directly as no dag is active.");
shutdownHandler.shutdown();
}
break;
case DAG_FINISHED:
DAGAppMasterEventDAGFinished finishEvt = (DAGAppMasterEventDAGFinished) event;
if (!isSession) {
setStateOnDAGCompletion();
LOG.info("Shutting down on completion of dag:" + finishEvt.getDAGId().toString());
shutdownHandler.shutdown();
} else {
LOG.info(
"DAG completed, dagId="
+ finishEvt.getDAGId().toString()
+ ", dagState="
+ finishEvt.getDAGState());
lastDAGCompletionTime = clock.getTime();
switch (finishEvt.getDAGState()) {
case SUCCEEDED:
successfulDAGs.incrementAndGet();
break;
case ERROR:
case FAILED:
failedDAGs.incrementAndGet();
break;
case KILLED:
killedDAGs.incrementAndGet();
break;
default:
LOG.fatal(
"Received a DAG Finished Event with state="
+ finishEvt.getDAGState()
+ ". Error. Shutting down.");
state = DAGAppMasterState.ERROR;
shutdownHandler.shutdown();
break;
}
if (!state.equals(DAGAppMasterState.ERROR)) {
if (!sessionStopped.get()) {
LOG.info("Waiting for next DAG to be submitted.");
state = DAGAppMasterState.IDLE;
} else {
LOG.info("Session shutting down now.");
state = DAGAppMasterState.SUCCEEDED;
shutdownHandler.shutdown();
}
}
}
break;
default:
throw new TezUncheckedException("AppMaster: No handler for event type: " + event.getType());
}
}