org.apache.tez.dag.api.client.DAGClient Java Examples

void runDAGAndVerify(DAG dag, DAGStatus.State finalState, int checkFailedAttempts, 
    String diagnostics) throws Exception {
  tezSession.waitTillReady();
  DAGClient dagClient = tezSession.submitDAG(dag);
  DAGStatus dagStatus = dagClient.getDAGStatus(null);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for dag to complete. Sleeping for 500ms."
        + " DAG name: " + dag.getName()
        + " DAG appContext: " + dagClient.getExecutionContext()
        + " Current state: " + dagStatus.getState());
    Thread.sleep(100);
    dagStatus = dagClient.getDAGStatus(null);
  }

  Assert.assertEquals(finalState, dagStatus.getState());
  
  if (checkFailedAttempts > 0) {
    Assert.assertEquals(checkFailedAttempts,
        dagStatus.getDAGProgress().getFailedTaskAttemptCount());
  }

  if (diagnostics != null) {
    Assert.assertNotNull(dagStatus.getDiagnostics());
    Assert.assertTrue(Joiner.on(":").join(dagStatus.getDiagnostics()).contains(diagnostics));
  }
}
 

void runDAGAndVerify(DAG dag, DAGStatus.State finalState,
                     TezClient session) throws Exception {
  session.waitTillReady();
  DAGClient dagClient = session.submitDAG(dag);
  DAGStatus dagStatus = dagClient.getDAGStatus(null);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for dag to complete. Sleeping for 500ms."
        + " DAG name: " + dag.getName()
        + " DAG appContext: " + dagClient.getExecutionContext()
        + " Current state: " + dagStatus.getState());
    Thread.sleep(100);
    dagStatus = dagClient.getDAGStatus(null);
  }

  Assert.assertEquals(finalState, dagStatus.getState());
}
 

@Test(timeout = 20000)
public void testSelfKillReported() throws IOException, TezException, InterruptedException {

  TezClient tezClient = getTezClient("testSelfKillReported");
  DAGClient dagClient = null;

  try {
    FailingProcessor.configureForKilled(10);
    DAG dag = DAG.create("testSelfKillReported").addVertex(
        Vertex
            .create(VERTEX_NAME, ProcessorDescriptor.create(FailingProcessor.class.getName()), 1));

    dagClient = tezClient.submitDAG(dag);
    dagClient.waitForCompletion();
    assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
    assertEquals(10, dagClient.getVertexStatus(VERTEX_NAME, null).getProgress().getKilledTaskAttemptCount());
  } finally {
    if (dagClient != null) {
      dagClient.close();
    }
    tezClient.stop();
  }
}
 

@Test(timeout = 20000)
public void testNonFatalErrorReported() throws IOException, TezException, InterruptedException {

  TezClient tezClient = getTezClient("testNonFatalErrorReported");
  DAGClient dagClient = null;

  try {
    FailingProcessor.configureForNonFatalFail();
    DAG dag = DAG.create("testNonFatalErrorReported").addVertex(
        Vertex
            .create(VERTEX_NAME, ProcessorDescriptor.create(FailingProcessor.class.getName()), 1));

    dagClient = tezClient.submitDAG(dag);
    dagClient.waitForCompletion();
    assertEquals(DAGStatus.State.FAILED, dagClient.getDAGStatus(null).getState());
    assertEquals(4, dagClient.getVertexStatus(VERTEX_NAME, null).getProgress().getFailedTaskAttemptCount());
  } finally {
    if (dagClient != null) {
      dagClient.close();
    }
    tezClient.stop();
  }
}
 

@Test(timeout = 20000)
public void testFatalErrorReported() throws IOException, TezException, InterruptedException {

  TezClient tezClient = getTezClient("testFatalErrorReported");
  DAGClient dagClient = null;

  try {
    FailingProcessor.configureForFatalFail();
    DAG dag = DAG.create("testFatalErrorReportedDag").addVertex(
        Vertex
            .create(VERTEX_NAME, ProcessorDescriptor.create(FailingProcessor.class.getName()), 1));

    dagClient = tezClient.submitDAG(dag);
    dagClient.waitForCompletion();
    assertEquals(DAGStatus.State.FAILED, dagClient.getDAGStatus(null).getState());
    assertEquals(1, dagClient.getVertexStatus(VERTEX_NAME, null).getProgress().getFailedTaskAttemptCount());
  } finally {
    if (dagClient != null) {
      dagClient.close();
    }
    tezClient.stop();
  }
}
 

@Test(timeout = 20000)
public void testNoSysExitOnSuccessfulDAG() throws TezException, InterruptedException,
    IOException {
  TezConfiguration tezConf1 = createConf();
  // Run in non-session mode so that the AM terminates
  TezClient tezClient1 = TezClient.create("commonName", tezConf1, false);
  tezClient1.start();

  DAG dag1 = createSimpleDAG("dag1", SleepProcessor.class.getName());

  DAGClient dagClient1 = tezClient1.submitDAG(dag1);
  dagClient1.waitForCompletion();
  assertEquals(DAGStatus.State.SUCCEEDED, dagClient1.getDAGStatus(null).getState());

  // Sleep for more time than is required for the DAG to complete.
  Thread.sleep((long) (TezConstants.TEZ_DAG_SLEEP_TIME_BEFORE_EXIT * 1.5));

  dagClient1.close();
  tezClient1.stop();
}
 

@Test(timeout = 20000)
public void testNoSysExitOnFailinglDAG() throws TezException, InterruptedException,
    IOException {
  TezConfiguration tezConf1 = createConf();
  // Run in non-session mode so that the AM terminates
  TezClient tezClient1 = TezClient.create("commonName", tezConf1, false);
  tezClient1.start();

  DAG dag1 = createSimpleDAG("dag1", FailingProcessor.class.getName());

  DAGClient dagClient1 = tezClient1.submitDAG(dag1);
  dagClient1.waitForCompletion();
  assertEquals(DAGStatus.State.FAILED, dagClient1.getDAGStatus(null).getState());

  // Sleep for more time than is required for the DAG to complete.
  Thread.sleep((long) (TezConstants.TEZ_DAG_SLEEP_TIME_BEFORE_EXIT * 1.5));

  dagClient1.close();
  tezClient1.stop();
}
 

/**
 * verify the diagnostics in DAGStatus is correct in session mode, using local
 * mode for fast speed
 * 
 * @throws Exception
 * 
 */
@Test(timeout = 600000)
public void testExceptionPropagationSession() throws Exception {
  try {
    startSessionClient();
    for (ExceptionLocation exLocation : ExceptionLocation.values()) {
      LOG.info("Session mode, Test for Exception from:" + exLocation.name());
      DAG dag = createDAG(exLocation);
      DAGClient dagClient = tezSession.submitDAG(dag);
      DAGStatus dagStatus = dagClient.waitForCompletion();
      String diagnostics = StringUtils.join(dagStatus.getDiagnostics(), ",");
      LOG.info("Diagnostics:" + diagnostics);
      if (exLocation == ExceptionLocation.PROCESSOR_COUNTER_EXCEEDED) {
        assertTrue(diagnostics.contains("Too many counters"));
      } else {
        assertTrue(diagnostics.contains(exLocation.name()));
      }
    }
  } finally {
    stopSessionClient();
  }
}
 

private void runDAG(DAG dag, DAGStatus.State finalState) throws Exception {
  tezSession.waitTillReady();
  LOG.info("ABC Running DAG name: " + dag.getName());
  DAGClient dagClient = tezSession.submitDAG(dag);
  DAGStatus dagStatus = dagClient.getDAGStatus(null);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for dag to complete. Sleeping for 500ms."
        + " DAG name: " + dag.getName()
        + " DAG appContext: " + dagClient.getExecutionContext()
        + " Current state: " + dagStatus.getState());
    Thread.sleep(100);
    dagStatus = dagClient.getDAGStatus(null);
  }

  Assert.assertEquals(finalState, dagStatus.getState());
}
 

private void testMemory(DAG dag, boolean sendDMEvents) throws Exception {
  StopWatch stopwatch = new StopWatch();
  stopwatch.start();
  TezConfiguration tezconf = new TezConfiguration(defaultConf);

  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null,
      null, false, false, numThreads, 1000);
  tezClient.start();
  
  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  mockApp.eventsDelegate = new TestMockDAGAppMaster.TestEventsDelegate();
  mockApp.doSleep = false;
  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  mockLauncher.startScheduling(true);
  DAGStatus status = dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, status.getState());
  checkMemory(dag.getName(), mockApp);
  stopwatch.stop();
  System.out.println("Time taken(ms): " + stopwatch.now(TimeUnit.MILLISECONDS));
  tezClient.stop();
}
 

@Ignore
@Test (timeout = 60000)
public void testTaskEventsProcessingSpeed() throws Exception {
  Logger.getRootLogger().setLevel(Level.WARN);
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  tezconf.setBoolean(TezConfiguration.TEZ_AM_USE_CONCURRENT_DISPATCHER, true);
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null,
      null, false, false, 30, 1000);
  tezClient.start();

  final String vAName = "A";
  
  DAG dag = DAG.create("testTaskEventsProcessingSpeed");
  Vertex vA = Vertex.create(vAName, ProcessorDescriptor.create("Proc.class"), 50000);
  dag.addVertex(vA);

  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  mockApp.doSleep = false;
  DAGClient dagClient = tezClient.submitDAG(dag);
  DAGStatus status = dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, status.getState());
  tezClient.stop();
}
 

void runDAGAndVerify(DAG dag, DAGStatus.State finalState,
                     TezClient session) throws Exception {
  session.waitTillReady();
  DAGClient dagClient = session.submitDAG(dag);
  DAGStatus dagStatus = dagClient.getDAGStatus(null);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for dag to complete. Sleeping for 500ms."
        + " DAG name: " + dag.getName()
        + " DAG appId: " + dagClient.getApplicationId()
        + " Current state: " + dagStatus.getState());
    Thread.sleep(100);
    dagStatus = dagClient.getDAGStatus(null);
  }

  Assert.assertEquals(finalState, dagStatus.getState());
}
 

@Test(timeout = 10000)
public void testMultipleClientsWithoutSession() throws TezException, InterruptedException,
    IOException {
  TezConfiguration tezConf1 = createConf();
  TezClient tezClient1 = TezClient.create("commonName", tezConf1, false);
  tezClient1.start();

  DAG dag1 = createSimpleDAG("dag1", SleepProcessor.class.getName());

  DAGClient dagClient1 = tezClient1.submitDAG(dag1);
  dagClient1.waitForCompletion();
  assertEquals(DAGStatus.State.SUCCEEDED, dagClient1.getDAGStatus(null).getState());

  dagClient1.close();
  tezClient1.stop();


  TezConfiguration tezConf2 = createConf();
  DAG dag2 = createSimpleDAG("dag2", SleepProcessor.class.getName());
  TezClient tezClient2 = TezClient.create("commonName", tezConf2, false);
  tezClient2.start();
  DAGClient dagClient2 = tezClient2.submitDAG(dag2);
  dagClient2.waitForCompletion();
  assertEquals(DAGStatus.State.SUCCEEDED, dagClient2.getDAGStatus(null).getState());
  assertFalse(dagClient1.getExecutionContext().equals(dagClient2.getExecutionContext()));
  dagClient2.close();
  tezClient2.stop();
}
 

@Test(timeout = 60000)
public void testVertexFailuresMaxPercent() throws TezException, InterruptedException, IOException {

  TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
  tezConf.set(TezConfiguration.TEZ_VERTEX_FAILURES_MAXPERCENT, "50.0f");
  tezConf.setInt(TezConfiguration.TEZ_AM_TASK_MAX_FAILED_ATTEMPTS, 1);
  TezClient tezClient = TezClient.create("TestVertexFailuresMaxPercent", tezConf);
  tezClient.start();

  try {
    DAG dag = DAG.create("TestVertexFailuresMaxPercent");
    Vertex vertex1 = Vertex.create("Parent", ProcessorDescriptor.create(
        FailingAttemptProcessor.class.getName()), 2);
    Vertex vertex2 = Vertex.create("Child", ProcessorDescriptor.create(FailingAttemptProcessor.class.getName()), 2);

    OrderedPartitionedKVEdgeConfig edgeConfig = OrderedPartitionedKVEdgeConfig
        .newBuilder(Text.class.getName(), IntWritable.class.getName(),
            HashPartitioner.class.getName())
        .setFromConfiguration(tezConf)
        .build();
    dag.addVertex(vertex1)
        .addVertex(vertex2)
        .addEdge(Edge.create(vertex1, vertex2, edgeConfig.createDefaultEdgeProperty()));

    DAGClient dagClient = tezClient.submitDAG(dag);
    dagClient.waitForCompletion();
    Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  } finally {
    tezClient.stop();
  }
}
 

@Test(timeout = 60000)
public void testInputInitializerEvents() throws TezException, InterruptedException, IOException {

  TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
  TezClient tezClient = TezClient.create("TestInputInitializerEvents", tezConf);
  tezClient.start();

  try {
    DAG dag = DAG.create("TestInputInitializerEvents");
    Vertex vertex1 = Vertex.create(VERTEX_WITH_INITIALIZER_NAME, ProcessorDescriptor.create(
        SleepProcessor.class.getName())
        .setUserPayload(new SleepProcessor.SleepProcessorConfig(1).toUserPayload()), 1)
        .addDataSource(INPUT1_NAME,
            DataSourceDescriptor
                .create(InputDescriptor.create(MultiAttemptDAG.NoOpInput.class.getName()),
                    InputInitializerDescriptor.create(InputInitializerForTest.class.getName()),
                    null));
    Vertex vertex2 = Vertex.create(EVENT_GENERATING_VERTEX_NAME,
        ProcessorDescriptor.create(InputInitializerEventGeneratingProcessor.class.getName()), 5);

    dag.addVertex(vertex1).addVertex(vertex2);

    DAGClient dagClient = tezClient.submitDAG(dag);
    dagClient.waitForCompletion();
    Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  } finally {
    tezClient.stop();
  }
}
 

private void runAndVerifyForNonFatalErrors(TezClient tezClient, String componentName,
                                           Vertex.VertexExecutionContext lhsContext) throws
    TezException,
    InterruptedException, IOException {
  LOG.info("Running JoinValidate with componentName reportNonFatalException");
  JoinValidateConfigured joinValidate =
      new JoinValidateConfigured(EXECUTION_CONTEXT_DEFAULT, lhsContext,
          EXECUTION_CONTEXT_EXT_SERVICE_PUSH,
          EXECUTION_CONTEXT_EXT_SERVICE_PUSH, componentName);

  DAG dag = joinValidate
      .createDag(new TezConfiguration(extServiceTestHelper.getConfForJobs()),
          HASH_JOIN_EXPECTED_RESULT_PATH,
          HASH_JOIN_OUTPUT_PATH, 3);

  DAGClient dagClient = tezClient.submitDAG(dag);

  DAGStatus dagStatus =
      dagClient.waitForCompletionWithStatusUpdates(Sets.newHashSet(StatusGetOpts.GET_COUNTERS));
  assertEquals(DAGStatus.State.FAILED, dagStatus.getState());

  boolean foundDiag = false;
  for (String diag : dagStatus.getDiagnostics()) {
    if (diag.contains(ErrorPluginConfiguration.REPORT_NONFATAL_ERROR_MESSAGE) &&
        diag.contains(ServicePluginErrorDefaults.SERVICE_UNAVAILABLE.name())) {
      foundDiag = true;
      break;
    }
  }
  assertTrue(foundDiag);
}
 

void runDAGAndVerify(DAG dag, DAGStatus.State finalState) throws Exception {
  tezSession.waitTillReady();
  DAGClient dagClient = tezSession.submitDAG(dag);
  DAGStatus dagStatus = dagClient.getDAGStatus(null);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for dag to complete. Sleeping for 500ms."
        + " DAG name: " + dag.getName()
        + " DAG appId: " + dagClient.getApplicationId()
        + " Current state: " + dagStatus.getState());
    Thread.sleep(100);
    dagStatus = dagClient.getDAGStatus(null);
  }
  
  Assert.assertEquals(finalState, dagStatus.getState());
}
 

@Test (timeout=50000)
public void testDisabledACls() throws Exception {
  TezClient tezSession = null;
  try {
    SleepProcessorConfig spConf = new SleepProcessorConfig(1);

    DAG dag = DAG.create("TezSleepProcessor");
    Vertex vertex = Vertex.create("SleepVertex", ProcessorDescriptor.create(
            SleepProcessor.class.getName()).setUserPayload(spConf.toUserPayload()), 1,
        Resource.newInstance(256, 1));
    dag.addVertex(vertex);

    TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_ALLOW_DISABLED_TIMELINE_DOMAINS, true);
    tezConf.set(TezConfiguration.TEZ_HISTORY_LOGGING_SERVICE_CLASS,
        ATSHistoryLoggingService.class.getName());
    Path remoteStagingDir = remoteFs.makeQualified(new Path("/tmp", String.valueOf(random
        .nextInt(100000))));
    remoteFs.mkdirs(remoteStagingDir);
    tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, remoteStagingDir.toString());

    tezSession = TezClient.create("TezSleepProcessor", tezConf, true);
    tezSession.start();

    DAGClient dagClient = tezSession.submitDAG(dag);

    DAGStatus dagStatus = dagClient.getDAGStatus(null);
    while (!dagStatus.isCompleted()) {
      LOG.info("Waiting for job to complete. Sleeping for 500ms." + " Current state: "
          + dagStatus.getState());
      Thread.sleep(500l);
      dagStatus = dagClient.getDAGStatus(null);
    }
    Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagStatus.getState());
  } finally {
    if (tezSession != null) {
      tezSession.stop();
    }
  }
}
 

@Test(timeout = 60000)
public void testSleepJob() throws TezException, IOException, InterruptedException {
  SleepProcessorConfig spConf = new SleepProcessorConfig(1);

  DAG dag = new DAG("TezSleepProcessor");
  Vertex vertex = new Vertex("SleepVertex", new ProcessorDescriptor(
      SleepProcessor.class.getName()).setUserPayload(spConf.toUserPayload()), 1,
      Resource.newInstance(1024, 1));
  dag.addVertex(vertex);

  TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
  Path remoteStagingDir = remoteFs.makeQualified(new Path("/tmp", String.valueOf(random
      .nextInt(100000))));
  remoteFs.mkdirs(remoteStagingDir);
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, remoteStagingDir.toString());

  TezClient tezSession = new TezClient("TezSleepProcessor", tezConf, false);
  tezSession.start();

  DAGClient dagClient = tezSession.submitDAG(dag);

  DAGStatus dagStatus = dagClient.getDAGStatus(null);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for job to complete. Sleeping for 500ms." + " Current state: "
        + dagStatus.getState());
    Thread.sleep(500l);
    dagStatus = dagClient.getDAGStatus(null);
  }
  dagStatus = dagClient.getDAGStatus(Sets.newHashSet(StatusGetOpts.GET_COUNTERS));

  assertEquals(DAGStatus.State.SUCCEEDED, dagStatus.getState());
  assertNotNull(dagStatus.getDAGCounters());
  assertNotNull(dagStatus.getDAGCounters().getGroup(FileSystemCounter.class.getName()));
  assertNotNull(dagStatus.getDAGCounters().findCounter(TaskCounter.GC_TIME_MILLIS));
  ExampleDriver.printDAGStatus(dagClient, new String[] { "SleepVertex" }, true, true);
  tezSession.stop();
}
 

@Test(timeout = 5000)
public void testDAGFinishedRecoveryError() throws Exception {
  TezConfiguration tezconf = new TezConfiguration(defaultConf);

  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();

  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  mockApp.recoveryFatalError = true;
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(true);

  DAG dag = DAG.create("test");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 5);
  dag.addVertex(vA);

  DAGClient dagClient = tezClient.submitDAG(dag);
  dagClient.waitForCompletion();
  while(!mockApp.getShutdownHandler().wasShutdownInvoked()) {
    Thread.sleep(100);
  }
  Assert.assertEquals(DAGState.SUCCEEDED, mockApp.getContext().getCurrentDAG().getState());
  Assert.assertEquals(DAGAppMasterState.FAILED, mockApp.getState());
  Assert.assertTrue(StringUtils.join(mockApp.getDiagnostics(),",")
      .contains("Recovery had a fatal error, shutting down session after" +
            " DAG completion"));
}
 

TezCounters runDAGAndVerify(DAG dag, DAGStatus.State finalState) throws Exception {
  tezSession.waitTillReady();
  DAGClient dagClient = tezSession.submitDAG(dag);
  DAGStatus dagStatus =
      dagClient.waitForCompletionWithStatusUpdates(EnumSet
          .of(StatusGetOpts.GET_COUNTERS));
  Assert.assertEquals(finalState, dagStatus.getState());
  return dagStatus.getDAGCounters();
}
 

@Test(timeout = 30000)
public void testMultipleClientsWithSession() throws TezException, InterruptedException,
    IOException {
  TezConfiguration tezConf1 = createConf();
  TezClient tezClient1 = TezClient.create("commonName", tezConf1, true);
  tezClient1.start();

  DAG dag1 = createSimpleDAG("dag1", SleepProcessor.class.getName());

  DAGClient dagClient1 = tezClient1.submitDAG(dag1);
  dagClient1.waitForCompletion();
  assertEquals(DAGStatus.State.SUCCEEDED, dagClient1.getDAGStatus(null).getState());

  dagClient1.close();
  tezClient1.stop();

  TezConfiguration tezConf2 = createConf();
  DAG dag2 = createSimpleDAG("dag2", SleepProcessor.class.getName());
  TezClient tezClient2 = TezClient.create("commonName", tezConf2, true);
  tezClient2.start();
  DAGClient dagClient2 = tezClient2.submitDAG(dag2);
  dagClient2.waitForCompletion();
  assertEquals(DAGStatus.State.SUCCEEDED, dagClient2.getDAGStatus(null).getState());
  assertFalse(dagClient1.getExecutionContext().equals(dagClient2.getExecutionContext()));
  dagClient2.close();
  tezClient2.stop();
}
 

boolean run(Configuration conf, String className, String confFilePath) throws Exception {
  this.conf = conf;
  setup();
  
  try {
    tezSession.waitTillReady();
    
    DAG dag = getDAG(className, confFilePath);
    
    DAGClient dagClient = tezSession.submitDAG(dag);
    DAGStatus dagStatus = dagClient.getDAGStatus(null);
    while (!dagStatus.isCompleted()) {
      System.out.println("Waiting for dag to complete. Sleeping for 500ms."
          + " DAG name: " + dag.getName()
          + " DAG appId: " + dagClient.getApplicationId()
          + " Current state: " + dagStatus.getState());
      Thread.sleep(500);
      dagStatus = dagClient.getDAGStatus(null);
    }
    
    if (dagStatus.getState() == DAGStatus.State.SUCCEEDED) {
      return true;
    }
    
  } finally {
    tearDown();
  }
  
  return false;
}
 

void runDAGAndVerify(DAG dag, DAGStatus.State finalState) throws Exception {
  tezSession.waitTillReady();
  DAGClient dagClient = tezSession.submitDAG(dag);
  DAGStatus dagStatus = dagClient.getDAGStatus(null, 10);
  while (!dagStatus.isCompleted()) {
    LOG.info("Waiting for dag to complete. Sleeping for 500ms."
        + " DAG name: " + dag.getName()
        + " DAG appContext: " + dagClient.getExecutionContext()
        + " Current state: " + dagStatus.getState());
    Thread.sleep(100);
    dagStatus = dagClient.getDAGStatus(null);
  }

  Assert.assertEquals(finalState, dagStatus.getState());
}
 

@Test (timeout = 10000)
public void testMultipleSubmissions() throws Exception {
  Map<String, LocalResource> lrDAG = Maps.newHashMap();
  String lrName1 = "LR1";
  lrDAG.put(lrName1, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test"),
      LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));
  Map<String, LocalResource> lrVertex = Maps.newHashMap();
  String lrName2 = "LR2";
  lrVertex.put(lrName2, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test1"),
      LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));

  DAG dag = DAG.create("test").addTaskLocalFiles(lrDAG);
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 5).addTaskLocalFiles(lrVertex);
  dag.addVertex(vA);

  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();
  DAGClient dagClient = tezClient.submitDAG(dag);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  tezClient.stop();
  
  // submit the same DAG again to verify it can be done.
  tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();
  dagClient = tezClient.submitDAG(dag);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  tezClient.stop();
}
 

/**
 * Test basic speculation not useful.
 *
 * @throws Exception the exception
 */
@Retry
@Test (timeout=10000)
public void testBasicSpeculationNotUseful() throws Exception {
  DAG dag = DAG.create("test");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 5);
  dag.addVertex(vA);

  MockTezClient tezClient = createTezSession();
  
  DAGClient dagClient = tezClient.submitDAG(dag);
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  TezVertexID vertexId = TezVertexID.getInstance(dagImpl.getID(), 0);
  // original attempt is successful and speculative one is killed
  TezTaskAttemptID successTaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), 0);
  TezTaskAttemptID killedTaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), 1);

  mockLauncher.setStatusUpdatesForTask(successTaId, 100);
  mockLauncher.setStatusUpdatesForTask(killedTaId, 100);

  mockLauncher.startScheduling(true);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  Task task = dagImpl.getTask(killedTaId.getTaskID());
  Assert.assertEquals(2, task.getAttempts().size());
  Assert.assertEquals(successTaId, task.getSuccessfulAttempt().getID());
  TaskAttempt killedAttempt = task.getAttempt(killedTaId);
  Joiner.on(",").join(killedAttempt.getDiagnostics()).contains("Killed speculative attempt as");
  Assert.assertEquals(TaskAttemptTerminationCause.TERMINATED_INEFFECTIVE_SPECULATION, 
      killedAttempt.getTerminationCause());
  Assert.assertEquals(1, task.getCounters().findCounter(TaskCounter.NUM_SPECULATIONS)
      .getValue());
  Assert.assertEquals(1, dagImpl.getAllCounters().findCounter(TaskCounter.NUM_SPECULATIONS)
      .getValue());
  org.apache.tez.dag.app.dag.Vertex v = dagImpl.getVertex(killedTaId.getTaskID().getVertexID());
  Assert.assertEquals(1, v.getAllCounters().findCounter(TaskCounter.NUM_SPECULATIONS)
      .getValue());
  tezClient.stop();
}
 

@Test (timeout = 5000)
public void testPreemptionWithoutSession() throws Exception {
  System.out.println("TestPreemptionWithoutSession");
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  tezconf.setInt(TezConfiguration.TEZ_AM_TASK_MAX_FAILED_ATTEMPTS, 0);
  AtomicBoolean mockAppLauncherGoFlag = new AtomicBoolean(false);
  MockTezClient tezClient = new MockTezClient("testPreemption", tezconf, false, null, null,
      null, mockAppLauncherGoFlag, false, false, 2, 2);
  tezClient.start();
  
  DAGClient dagClient = tezClient.submitDAG(createDAG(DataMovementType.SCATTER_GATHER));
  // now the MockApp has been started. sync with it to get the launcher
  syncWithMockAppLauncher(false, mockAppLauncherGoFlag, tezClient);

  DAGImpl dagImpl;
  do {
    Thread.sleep(100); // usually needs to sleep 2-3 times
  } while ((dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG()) == null);

  int vertexIndex = 0;
  int upToTaskVersion = 3;
  TezVertexID vertexId = TezVertexID.getInstance(dagImpl.getID(), vertexIndex);
  TezTaskAttemptID taId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), 0);

  mockLauncher.preemptContainerForTask(taId.getTaskID(), upToTaskVersion);
  mockLauncher.startScheduling(true);
  
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());

  for (int i=0; i<=upToTaskVersion; ++i) {
    TezTaskAttemptID testTaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), i);      
    TaskAttemptImpl taImpl = dagImpl.getTaskAttempt(testTaId);
    Assert.assertEquals(TaskAttemptStateInternal.KILLED, taImpl.getInternalState());
  }
  
  tezClient.stop();
}
 

void testPreemptionJob(MockTezClient tezClient, DAG dag, int vertexIndex,
    int upToTaskVersion, String info) throws Exception {
  System.out.println("TestPreemption - Running - " + info);
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  tezconf.setInt(TezConfiguration.TEZ_AM_TASK_MAX_FAILED_ATTEMPTS, 0);
  
  mockLauncher.startScheduling(false); // turn off scheduling to block DAG before submitting it
  DAGClient dagClient = tezClient.submitDAG(dag);
  
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  TezVertexID vertexId = TezVertexID.getInstance(dagImpl.getID(), vertexIndex);
  TezTaskAttemptID taId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), 0);

  mockLauncher.preemptContainerForTask(taId.getTaskID(), upToTaskVersion);
  mockLauncher.startScheduling(true);
  
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  
  for (int i=0; i<=upToTaskVersion; ++i) {
    TezTaskAttemptID testTaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), i);      
    TaskAttemptImpl taImpl = dagImpl.getTaskAttempt(testTaId);
    Assert.assertEquals(TaskAttemptStateInternal.KILLED, taImpl.getInternalState());
    Assert.assertEquals(TaskAttemptTerminationCause.EXTERNAL_PREEMPTION, taImpl.getTerminationCause());
  }
  
  System.out.println("TestPreemption - Done running - " + info);
}
 

@Test (timeout = 5000)
public void testLocalResourceSetup() throws Exception {
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();
  
  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  
  Map<String, LocalResource> lrDAG = Maps.newHashMap();
  String lrName1 = "LR1";
  lrDAG.put(lrName1, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test"),
      LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));
  Map<String, LocalResource> lrVertex = Maps.newHashMap();
  String lrName2 = "LR2";
  lrVertex.put(lrName2, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test1"),
      LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));

  DAG dag = DAG.create("testLocalResourceSetup").addTaskLocalFiles(lrDAG);
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 5).addTaskLocalFiles(lrVertex);
  dag.addVertex(vA);

  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  ContainerData cData = mockLauncher.getContainers().values().iterator().next();
  ContainerLaunchContext launchContext = cData.launchContext;
  Map<String, LocalResource> taskLR = launchContext.getLocalResources();
  // verify tasks are launched with both DAG and task resources.
  Assert.assertTrue(taskLR.containsKey(lrName1));
  Assert.assertTrue(taskLR.containsKey(lrName2));
  
  mockLauncher.startScheduling(true);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  tezClient.stop();
}
 

@Test (timeout = 5000)
public void testInternalPreemption() throws Exception {
  TezConfiguration tezconf = new TezConfiguration(defaultConf);
  
  MockTezClient tezClient = new MockTezClient("testMockAM", tezconf, true, null, null, null, null);
  tezClient.start();
  
  MockDAGAppMaster mockApp = tezClient.getLocalClient().getMockApp();
  MockContainerLauncher mockLauncher = mockApp.getContainerLauncher();
  mockLauncher.startScheduling(false);
  // there is only 1 task whose first attempt will be preempted
  DAG dag = DAG.create("testInternalPreemption");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), 1);
  dag.addVertex(vA);

  DAGClient dagClient = tezClient.submitDAG(dag);
  mockLauncher.waitTillContainersLaunched();
  ContainerData cData = mockLauncher.getContainers().values().iterator().next();
  DAGImpl dagImpl = (DAGImpl) mockApp.getContext().getCurrentDAG();
  mockApp.getTaskSchedulerManager().preemptContainer(0, cData.cId);
  
  mockLauncher.startScheduling(true);
  dagClient.waitForCompletion();
  Assert.assertEquals(DAGStatus.State.SUCCEEDED, dagClient.getDAGStatus(null).getState());
  TezVertexID vertexId = TezVertexID.getInstance(dagImpl.getID(), 0);
  TezTaskAttemptID killedTaId = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, 0), 0);
  TaskAttempt killedTa = dagImpl.getVertex(vA.getName()).getTask(0).getAttempt(killedTaId);
  //Refer to TEZ-3950
  Assert.assertTrue(killedTa.getState().equals(TaskAttemptState.KILLED) || killedTa.getState().equals(TaskAttemptState.FAILED));
  tezClient.stop();
}