@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);
 }
示例#10
0
  @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());
   }
 }