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