Java Code Examples for org.apache.tez.dag.api.TezConfiguration#setInt()

@Before
public void setup()  throws Exception {
  LOG.info("Starting session");
  Path remoteStagingDir = remoteFs.makeQualified(new Path(TEST_ROOT_DIR, String
      .valueOf(new Random().nextInt(100000))));
  TezClientUtils.ensureStagingDirExists(conf, remoteStagingDir);

  TezConfiguration tezConf = new TezConfiguration(miniTezCluster.getConfig());
  tezConf.setInt(TezConfiguration.DAG_RECOVERY_MAX_UNFLUSHED_EVENTS, 0);
  tezConf.set(TezConfiguration.TEZ_AM_LOG_LEVEL, "DEBUG");
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR,
      remoteStagingDir.toString());
  tezConf.setBoolean(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_ENABLED, false);
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.setInt(TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, 500);
  tezConf.set(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS, " -Xmx256m");
  tezConf.setBoolean(TezConfiguration.TEZ_AM_SESSION_MODE, true);

  tezSession = new TezClient("TestDAGRecovery", tezConf);
  tezSession.start();
}
 

private void startSessionClient() throws Exception {
  LOG.info("Starting session");
  tezConf = new TezConfiguration();
  tezConf.setInt(TezConfiguration.DAG_RECOVERY_MAX_UNFLUSHED_EVENTS, 0);
  tezConf
      .setBoolean(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_ENABLED, false);
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.setBoolean(TezConfiguration.TEZ_AM_ONE_TO_ONE_ROUTING_USE_ON_DEMAND_ROUTING, true);
  tezConf.setInt(TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, 500);
  tezConf.set(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS, " -Xmx256m");
  tezConf.setBoolean(TezConfiguration.TEZ_AM_SESSION_MODE, true);
  // for local mode
  tezConf.setBoolean(TezConfiguration.TEZ_LOCAL_MODE, true);
  tezConf.set("fs.defaultFS", "file:///");
  tezConf.setBoolean(
      TezRuntimeConfiguration.TEZ_RUNTIME_OPTIMIZE_LOCAL_FETCH, true);

  tezSession = TezClient.create("TestExceptionPropagation", tezConf);
  tezSession.start();
}
 

@Before
public void setup()  throws Exception {
  LOG.info("Starting session");
  Path remoteStagingDir = remoteFs.makeQualified(new Path(TEST_ROOT_DIR, String
      .valueOf(new Random().nextInt(100000))));
  TezClientUtils.ensureStagingDirExists(conf, remoteStagingDir);

  tezConf = new TezConfiguration(miniTezCluster.getConfig());
  tezConf.setInt(TezConfiguration.DAG_RECOVERY_MAX_UNFLUSHED_EVENTS, 0);
  tezConf.set(TezConfiguration.TEZ_AM_LOG_LEVEL, "INFO");
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR,
      remoteStagingDir.toString());
  tezConf.setBoolean(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_ENABLED, false);
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.setInt(TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, 500);
  tezConf.set(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS, " -Xmx256m");
  tezConf.setBoolean(TezConfiguration.TEZ_AM_SESSION_MODE, true);
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_SCRATCH_DATA_AUTO_DELETE, "false");
  tezConf.setInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_KEY,0);
  tezConf.setInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_ON_SOCKET_TIMEOUTS_KEY, 0);
  tezConf.setInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_TIMEOUT_KEY,1000);

  tezSession = TezClient.create("TestDAGRecovery", tezConf);
  tezSession.start();
}
 

@Test(timeout = 20000)
public void testAMHeartbeatFailOnGetAMProxy_AppFailed() throws Exception {
  int amHeartBeatTimeoutSecs = 3;
  TezConfiguration conf = new TezConfiguration();
  conf.setInt(TezConfiguration.TEZ_AM_CLIENT_HEARTBEAT_TIMEOUT_SECS, amHeartBeatTimeoutSecs);

  final TezClientForTest client = configureAndCreateTezClient(conf);
  client.callRealGetSessionAMProxy = true;
  client.start();

  when(client.mockYarnClient.getApplicationReport(client.mockAppId).getYarnApplicationState())
    .thenReturn(YarnApplicationState.FAILED);
  Thread.sleep(2 * amHeartBeatTimeoutSecs * 1000);
  assertTrue(client.getAMKeepAliveService().isTerminated());
}
 

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

MockTezClient createTezSession() throws Exception {
  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, true, null, null,
      null, mockAppLauncherGoFlag);
  tezClient.start();
  syncWithMockAppLauncher(false, mockAppLauncherGoFlag, tezClient);
  return tezClient;
}
 

@Test(timeout = 20000)
public void testAMHeartbeatFailOnGetAMProxy_AppNotStarted() throws Exception {
  int amHeartBeatTimeoutSecs = 3;
  TezConfiguration conf = new TezConfiguration();
  conf.setInt(TezConfiguration.TEZ_AM_CLIENT_HEARTBEAT_TIMEOUT_SECS, amHeartBeatTimeoutSecs);

  final TezClientForTest client = configureAndCreateTezClient(conf);
  client.callRealGetSessionAMProxy = true;
  client.start();

  when(client.mockYarnClient.getApplicationReport(client.mockAppId).getYarnApplicationState())
      .thenReturn(YarnApplicationState.ACCEPTED);
  Thread.sleep(2 * amHeartBeatTimeoutSecs * 1000);
  assertFalse(client.getAMKeepAliveService().isTerminated());
}
 

private TezConfiguration createSessionConfig(Path remoteStagingDir) {
  TezConfiguration tezConf = new TezConfiguration(miniTezCluster.getConfig());
  tezConf.setInt(TezConfiguration.DAG_RECOVERY_MAX_UNFLUSHED_EVENTS, 10);
  tezConf.set(TezConfiguration.TEZ_AM_LOG_LEVEL, "DEBUG");
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR,
      remoteStagingDir.toString());
  tezConf.setBoolean(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_ENABLED, false);
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.setInt(TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, 500);
  tezConf.set(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS, " -Xmx256m");
  tezConf.setBoolean(TezConfiguration.TEZ_AM_SESSION_MODE, true);
  return tezConf;
}
 

@Test(timeout = 60000)
public void testCartesianProduct() throws Exception {
  LOG.info("Running CartesianProduct Test");
  CartesianProduct job = new CartesianProduct();

  TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
  tezConf.setInt(CartesianProductVertexManager.TEZ_CARTESIAN_PRODUCT_MAX_PARALLELISM, 10);
  tezConf.setInt(CartesianProductVertexManager.TEZ_CARTESIAN_PRODUCT_MIN_OPS_PER_WORKER, 25);
  Assert.assertEquals("CartesianProduct failed", job.run(tezConf, null, null), 0);
}
 

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

@BeforeClass
public static void setup() throws Exception {
  LOG.info("Starting mini clusters");
  FileSystem remoteFs = null;
  try {
    conf.set(MiniDFSCluster.HDFS_MINIDFS_BASEDIR, TEST_ROOT_DIR);
    dfsCluster = new MiniDFSCluster.Builder(conf).numDataNodes(1)
        .format(true).racks(null).build();
    remoteFs = dfsCluster.getFileSystem();
  } catch (IOException io) {
    throw new RuntimeException("problem starting mini dfs cluster", io);
  }
  if (miniTezCluster == null) {
    miniTezCluster = new MiniTezCluster(TestFaultTolerance.class.getName(),
        3, 1, 1);
    Configuration miniTezconf = new Configuration(conf);
    miniTezconf.set("fs.defaultFS", remoteFs.getUri().toString()); // use HDFS
    miniTezCluster.init(miniTezconf);
    miniTezCluster.start();
    
    Path remoteStagingDir = remoteFs.makeQualified(new Path(TEST_ROOT_DIR, String
        .valueOf(new Random().nextInt(100000))));
    TezClientUtils.ensureStagingDirExists(conf, remoteStagingDir);
    
    tezConf = new TezConfiguration(miniTezCluster.getConfig());
    tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR,
        remoteStagingDir.toString());
    tezConf.setBoolean(TezConfiguration.TEZ_AM_NODE_BLACKLISTING_ENABLED, false);
    tezConf.setDouble(TezConfiguration.TEZ_TASK_MAX_ALLOWED_OUTPUT_FAILURES_FRACTION, 0.4);
    tezConf.setInt(TezConfiguration.TEZ_AM_MAX_ALLOWED_TIME_FOR_TASK_READ_ERROR_SEC, 3);
    tezConf.setInt(TezConfiguration.TEZ_TASK_AM_HEARTBEAT_INTERVAL_MS, 100);
    tezConf.setLong(TezConfiguration.TEZ_AM_SLEEP_TIME_BEFORE_EXIT_MILLIS, 500);

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

@Test
public void maxTasksAllocationsCannotBeExceeded() {

  final int MAX_TASKS = 4;
  TezConfiguration tezConf = new TezConfiguration();
  tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS);

  LocalContainerFactory containerFactory = new LocalContainerFactory(createMockAppContext());
  HashMap<Object, Container> taskAllocations = new LinkedHashMap<Object, Container>();
  PriorityBlockingQueue<TaskRequest> taskRequestQueue = new PriorityBlockingQueue<TaskRequest>();
  TaskSchedulerAppCallback appClientDelegate = mock(TaskSchedulerAppCallback.class);

  // Object under test
  AsyncDelegateRequestHandler requestHandler =
    new AsyncDelegateRequestHandler(taskRequestQueue,
        containerFactory,
        taskAllocations,
        appClientDelegate,
        tezConf);

  // Allocate up to max tasks
  for (int i = 0; i < MAX_TASKS; i++) {
    Priority priority = Priority.newInstance(20);
    requestHandler.addAllocateTaskRequest(new Long(i), null, priority, null);
    requestHandler.processRequest();
  }

  // Only MAX_TASKS number of tasks should have been allocated
  Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, taskAllocations.size());
  Assert.assertTrue("Another allocation should not fit", requestHandler.shouldWait());

  // Deallocate down to zero
  for (int i = 0; i < MAX_TASKS; i++) {
    requestHandler.addDeallocateTaskRequest(new Long(i));
    requestHandler.processRequest();
  }

  // All allocated tasks should have been removed
  Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size());
}
 

@Test(timeout = 5000)
public void maxTasksAllocationsCannotBeExceeded() {

  final int MAX_TASKS = 4;
  TezConfiguration tezConf = new TezConfiguration();
  tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS);

  ApplicationId appId = ApplicationId.newInstance(2000, 1);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);

  TaskSchedulerContext
      mockContext = TestTaskSchedulerHelpers.setupMockTaskSchedulerContext("", 0, "", true,
      appAttemptId, 1000l, null, new Configuration());

  LocalContainerFactory containerFactory = new LocalContainerFactory(appAttemptId, 1000);

  HashMap<Object, AllocatedTask> taskAllocations = new LinkedHashMap<>();
  LinkedBlockingQueue<SchedulerRequest> clientRequestQueue = new LinkedBlockingQueue<>();

  // Object under test
  AsyncDelegateRequestHandler requestHandler =
    new AsyncDelegateRequestHandler(clientRequestQueue,
        containerFactory,
        taskAllocations,
        mockContext,
        tezConf);

  // Allocate up to max tasks
  for (int i = 0; i < MAX_TASKS; i++) {
    Priority priority = Priority.newInstance(20);
    requestHandler.addAllocateTaskRequest(new Long(i), null, priority, null);
    requestHandler.dispatchRequest();
    requestHandler.allocateTask();
  }

  // Only MAX_TASKS number of tasks should have been allocated
  Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, taskAllocations.size());
  Assert.assertTrue("Another allocation should not fit", !requestHandler.shouldProcess());

  // Deallocate down to zero
  for (int i = 0; i < MAX_TASKS; i++) {
    requestHandler.addDeallocateTaskRequest(new Long(i));
    requestHandler.dispatchRequest();
  }

  // All allocated tasks should have been removed
  Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size());
}
 

@Test(timeout = 60000)
public void testSessionTimeout() throws Exception {
  Path stagingDirPath = new Path("/tmp/sessiontimeout-staging-dir");
  remoteFs.mkdirs(stagingDirPath);

  YarnClient yarnClient = YarnClient.createYarnClient();

  try {

    yarnClient.init(mrrTezCluster.getConfig());
    yarnClient.start();

    List<ApplicationReport> apps = yarnClient.getApplications();
    int appsBeforeCount = apps != null ? apps.size() : 0;

    TezConfiguration tezConf = new TezConfiguration(mrrTezCluster.getConfig());
    tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, stagingDirPath.toString());
    tezConf.setInt(TezConfiguration.TEZ_SESSION_AM_DAG_SUBMIT_TIMEOUT_SECS, 5);
    TezClient tezClient = TezClient.create("testSessionTimeout", tezConf, true);
    tezClient.start();

    ApplicationId appId = tezClient.getAppMasterApplicationId();

    apps = yarnClient.getApplications();
    int appsAfterCount = apps != null ? apps.size() : 0;

    // Running in session mode. So should only create 1 more app.
    Assert.assertEquals(appsBeforeCount + 1, appsAfterCount);

    ApplicationReport report;
    while (true) {
      report = yarnClient.getApplicationReport(appId);
      if (report.getYarnApplicationState() == YarnApplicationState.FINISHED
          || report.getYarnApplicationState() == YarnApplicationState.FAILED
          || report.getYarnApplicationState() == YarnApplicationState.KILLED) {
        break;
      }
      Thread.sleep(1000);
    }
    // Add a sleep because YARN is not consistent in terms of reporting uptodate diagnostics
    Thread.sleep(2000);
    report = yarnClient.getApplicationReport(appId);
    LOG.info("App Report for appId=" + appId
        + ", report=" + report);
    Assert.assertTrue("Actual diagnostics: " + report.getDiagnostics(),
        report.getDiagnostics().contains("Session timed out"));

  } finally {
    remoteFs.delete(stagingDirPath, true);
    if (yarnClient != null) {
      yarnClient.stop();
    }
  }
}
 

/**
 * In unpartitioned cartesian product, failure fraction should be #unique failure/#consumer that
 * depends on the src task. Here we test a 2x2 cartesian product and let 4th destination task fail.
 * The failure fraction limit is configured to be 0.25. So the failure fraction should be 1/2,
 * not 1/4.
 * @throws Exception
 */
@Test
public void testCartesianProduct() throws Exception {
  Configuration dagConf = new Configuration();
  dagConf.setDouble(TezConfiguration.TEZ_TASK_MAX_ALLOWED_OUTPUT_FAILURES_FRACTION, 0.25);
  DAG dag = DAG.create("dag");

  Configuration vertexConf = new Configuration();
  vertexConf.setInt(TestProcessor.getVertexConfName(
    TestProcessor.TEZ_FAILING_PROCESSOR_VERIFY_TASK_INDEX, "v3"), 3);
  vertexConf.setInt(TestProcessor.getVertexConfName(
    TestProcessor.TEZ_FAILING_PROCESSOR_VERIFY_VALUE, "v3"), 5);
  UserPayload vertexPayload = TezUtils.createUserPayloadFromConf(vertexConf);
  ProcessorDescriptor processorDescriptor =
    ProcessorDescriptor.create(TestProcessor.class.getName()).setUserPayload(vertexPayload);
  Vertex v1 = Vertex.create("v1", processorDescriptor, 2);
  Vertex v2 = Vertex.create("v2", processorDescriptor, 2);
  Vertex v3 = Vertex.create("v3", processorDescriptor);

  String[] sourceVertices = {"v1", "v2"};
  CartesianProductConfig cartesianProductConfig =
    new CartesianProductConfig(Arrays.asList(sourceVertices));
  TezConfiguration tezConf = new TezConfiguration();
  tezConf.setInt(CartesianProductVertexManager.TEZ_CARTESIAN_PRODUCT_NUM_PARTITIONS, 1);
  tezConf.setBoolean(CartesianProductVertexManager.TEZ_CARTESIAN_PRODUCT_ENABLE_GROUPING, false);
  UserPayload cartesianProductPayload =
    cartesianProductConfig.toUserPayload(tezConf);

  v3.setVertexManagerPlugin(
    VertexManagerPluginDescriptor.create(CartesianProductVertexManager.class.getName())
      .setUserPayload(cartesianProductPayload));

  EdgeManagerPluginDescriptor edgeManagerPluginDescriptor =
    EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName())
      .setUserPayload(cartesianProductPayload);

  Configuration inputConf = new Configuration();
  inputConf.setBoolean(TestInput.getVertexConfName(
    TestInput.TEZ_FAILING_INPUT_DO_FAIL, "v3"), true);
  inputConf.setInt(TestInput.getVertexConfName(
    TestInput.TEZ_FAILING_INPUT_FAILING_TASK_INDEX, "v3"), 3);
  inputConf.setInt(TestInput.getVertexConfName(
    TestInput.TEZ_FAILING_INPUT_FAILING_TASK_ATTEMPT, "v3"), 0);
  inputConf.setInt(TestInput.getVertexConfName(
    TestInput.TEZ_FAILING_INPUT_FAILING_INPUT_INDEX, "v3"), 0);
  inputConf.setInt(TestInput.getVertexConfName(
    TestInput.TEZ_FAILING_INPUT_FAILING_UPTO_INPUT_ATTEMPT, "v3"), 0);
  UserPayload inputPayload = TezUtils.createUserPayloadFromConf(inputConf);
  EdgeProperty edgeProperty =
    EdgeProperty.create(edgeManagerPluginDescriptor, DataMovementType.CUSTOM,
      DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, TestOutput.getOutputDesc(null),
      TestInput.getInputDesc(inputPayload));
  Edge e1 = Edge.create(v1, v3, edgeProperty);
  Edge e2 = Edge.create(v2, v3, edgeProperty);
  dag.addVertex(v1).addVertex(v2).addVertex(v3);
  dag.addEdge(e1).addEdge(e2);

  // run dag
  runDAGAndVerify(dag, DAGStatus.State.SUCCEEDED);
}
 

private void testHashJoinExample(SimpleShutdownCondition shutdownCondition,
    boolean enableAutoParallelism, boolean generateSplitInClient) throws Exception {
  HashJoinExample hashJoinExample = new HashJoinExample();
  TezConfiguration tezConf = new TezConfiguration(miniTezCluster.getConfig());
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.set(TezConfiguration.TEZ_AM_RECOVERY_SERVICE_CLASS,
      RecoveryServiceWithEventHandlingHook.class.getName());
  tezConf.set(
      RecoveryServiceWithEventHandlingHook.AM_RECOVERY_SERVICE_HOOK_CLASS,
      SimpleRecoveryEventHook.class.getName());
  tezConf.set(SimpleRecoveryEventHook.SIMPLE_SHUTDOWN_CONDITION,
      shutdownCondition.serialize());
  tezConf.setBoolean(
      ShuffleVertexManager.TEZ_SHUFFLE_VERTEX_MANAGER_ENABLE_AUTO_PARALLEL,
      enableAutoParallelism);
  tezConf.setBoolean(
      RecoveryService.TEZ_TEST_RECOVERY_DRAIN_EVENTS_WHEN_STOPPED, false);
  tezConf.setBoolean(
      TezConfiguration.TEZ_AM_STAGING_SCRATCH_DATA_AUTO_DELETE, false);
  tezConf.setInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_KEY,0);
  tezConf.setInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_ON_SOCKET_TIMEOUTS_KEY, 0);
  tezConf.setInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_TIMEOUT_KEY,1000);
  tezConf.set(TezConfiguration.TEZ_AM_LOG_LEVEL, "INFO;org.apache.tez=DEBUG");

  hashJoinExample.setConf(tezConf);
  Path stagingDirPath = new Path("/tmp/tez-staging-dir");
  Path inPath1 = new Path("/tmp/hashJoin/inPath1");
  Path inPath2 = new Path("/tmp/hashJoin/inPath2");
  Path outPath = new Path("/tmp/hashJoin/outPath");
  remoteFs.delete(outPath, true);
  remoteFs.mkdirs(inPath1);
  remoteFs.mkdirs(inPath2);
  remoteFs.mkdirs(stagingDirPath);

  Set<String> expectedResult = new HashSet<String>();

  FSDataOutputStream out1 = remoteFs.create(new Path(inPath1, "file"));
  FSDataOutputStream out2 = remoteFs.create(new Path(inPath2, "file"));
  BufferedWriter writer1 = new BufferedWriter(new OutputStreamWriter(out1));
  BufferedWriter writer2 = new BufferedWriter(new OutputStreamWriter(out2));
  for (int i = 0; i < 20; i++) {
    String term = "term" + i;
    writer1.write(term);
    writer1.newLine();
    if (i % 2 == 0) {
      writer2.write(term);
      writer2.newLine();
      expectedResult.add(term);
    }
  }
  writer1.close();
  writer2.close();
  out1.close();
  out2.close();

  String[] args = null;
  if (generateSplitInClient) {
    args = new String[]{
        "-D" + TezConfiguration.TEZ_AM_STAGING_DIR + "="
            + stagingDirPath.toString(),
        "-generateSplitInClient",
        inPath1.toString(), inPath2.toString(), "1", outPath.toString()};
  } else {
    args = new String[]{
        "-D" + TezConfiguration.TEZ_AM_STAGING_DIR + "="
            + stagingDirPath.toString(),
        inPath1.toString(), inPath2.toString(), "1", outPath.toString()};
  }
  assertEquals(0, hashJoinExample.run(args));

  FileStatus[] statuses = remoteFs.listStatus(outPath, new PathFilter() {
    public boolean accept(Path p) {
      String name = p.getName();
      return !name.startsWith("_") && !name.startsWith(".");
    }
  });
  assertEquals(1, statuses.length);
  FSDataInputStream inStream = remoteFs.open(statuses[0].getPath());
  BufferedReader reader = new BufferedReader(new InputStreamReader(inStream));
  String line;
  while ((line = reader.readLine()) != null) {
    assertTrue(expectedResult.remove(line));
  }
  reader.close();
  inStream.close();
  assertEquals(0, expectedResult.size());

  List<HistoryEvent> historyEventsOfAttempt1 = RecoveryParser
      .readRecoveryEvents(tezConf, hashJoinExample.getAppId(), 1);
  HistoryEvent lastEvent = historyEventsOfAttempt1
      .get(historyEventsOfAttempt1.size() - 1);
  assertEquals(shutdownCondition.getEvent().getEventType(),
      lastEvent.getEventType());
  assertTrue(shutdownCondition.match(lastEvent));
}
 

private void testOrderedWordCountMultipleRoundRecoverying(
        RecoveryServiceWithEventHandlingHook.MultipleRoundShutdownCondition shutdownCondition,
        boolean enableAutoParallelism, boolean generateSplitInClient) throws Exception {

  for (int i=0; i<shutdownCondition.size(); i++) {
    SimpleShutdownCondition condition = shutdownCondition.getSimpleShutdownCondition(i);
    LOG.info("ShutdownCondition:" + condition.getEventType()
            + ", event=" + condition.getEvent());
  }

  String inputDirStr = "/tmp/owc-input/";
  Path inputDir = new Path(inputDirStr);
  Path stagingDirPath = new Path("/tmp/owc-staging-dir");
  remoteFs.mkdirs(inputDir);
  remoteFs.mkdirs(stagingDirPath);
  TestTezJobs.generateOrderedWordCountInput(inputDir, remoteFs);

  String outputDirStr = "/tmp/owc-output/";
  Path outputDir = new Path(outputDirStr);

  TezConfiguration tezConf = new TezConfiguration(miniTezCluster.getConfig());
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.set(TezConfiguration.TEZ_AM_RECOVERY_SERVICE_CLASS,
          RecoveryServiceWithEventHandlingHook.class.getName());
  tezConf.set(
          RecoveryServiceWithEventHandlingHook.AM_RECOVERY_SERVICE_HOOK_CLASS,
          RecoveryServiceWithEventHandlingHook.MultipleRoundRecoveryEventHook.class.getName());
  tezConf.set(RecoveryServiceWithEventHandlingHook.MultipleRoundRecoveryEventHook.MULTIPLE_ROUND_SHUTDOWN_CONDITION,
          shutdownCondition.serialize());
  tezConf.setBoolean(
          ShuffleVertexManager.TEZ_SHUFFLE_VERTEX_MANAGER_ENABLE_AUTO_PARALLEL,
          enableAutoParallelism);
  tezConf.setBoolean(
          RecoveryService.TEZ_TEST_RECOVERY_DRAIN_EVENTS_WHEN_STOPPED, false);
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, stagingDirPath.toString());
  tezConf.setBoolean(
          TezConfiguration.TEZ_AM_STAGING_SCRATCH_DATA_AUTO_DELETE, false);
  OrderedWordCount job = new OrderedWordCount();
  if (generateSplitInClient) {
    Assert
            .assertTrue("OrderedWordCount failed", job.run(tezConf, new String[]{
                    "-generateSplitInClient", inputDirStr, outputDirStr, "5"}, null) == 0);
  } else {
    Assert
            .assertTrue("OrderedWordCount failed", job.run(tezConf, new String[]{
                    inputDirStr, outputDirStr, "5"}, null) == 0);
  }
  TestTezJobs.verifyOutput(outputDir, remoteFs);
}
 

private void testOrderedWordCount(SimpleShutdownCondition shutdownCondition,
    boolean enableAutoParallelism, boolean generateSplitInClient) throws Exception {
  LOG.info("shutdownCondition:" + shutdownCondition.getEventType()
      + ", event=" + shutdownCondition.getEvent());
  String inputDirStr = "/tmp/owc-input/";
  Path inputDir = new Path(inputDirStr);
  Path stagingDirPath = new Path("/tmp/owc-staging-dir");
  remoteFs.mkdirs(inputDir);
  remoteFs.mkdirs(stagingDirPath);
  TestTezJobs.generateOrderedWordCountInput(inputDir, remoteFs);

  String outputDirStr = "/tmp/owc-output/";
  Path outputDir = new Path(outputDirStr);

  TezConfiguration tezConf = new TezConfiguration(miniTezCluster.getConfig());
  tezConf.setInt(TezConfiguration.TEZ_AM_MAX_APP_ATTEMPTS, 4);
  tezConf.set(TezConfiguration.TEZ_AM_RECOVERY_SERVICE_CLASS,
      RecoveryServiceWithEventHandlingHook.class.getName());
  tezConf.set(
      RecoveryServiceWithEventHandlingHook.AM_RECOVERY_SERVICE_HOOK_CLASS,
      SimpleRecoveryEventHook.class.getName());
  tezConf.set(SimpleRecoveryEventHook.SIMPLE_SHUTDOWN_CONDITION,
      shutdownCondition.serialize());
  tezConf.setBoolean(
      ShuffleVertexManager.TEZ_SHUFFLE_VERTEX_MANAGER_ENABLE_AUTO_PARALLEL,
      enableAutoParallelism);
  tezConf.setBoolean(
      RecoveryService.TEZ_TEST_RECOVERY_DRAIN_EVENTS_WHEN_STOPPED, false);
  tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, stagingDirPath.toString());
  tezConf.setBoolean(
      TezConfiguration.TEZ_AM_STAGING_SCRATCH_DATA_AUTO_DELETE, false);
  tezConf.set(TezConfiguration.TEZ_AM_LOG_LEVEL, "INFO;org.apache.tez=DEBUG");
  OrderedWordCount job = new OrderedWordCount();
  if (generateSplitInClient) {
    Assert
        .assertTrue("OrderedWordCount failed", job.run(tezConf, new String[]{
            "-generateSplitInClient", inputDirStr, outputDirStr, "5"}, null) == 0);
  } else {
    Assert
        .assertTrue("OrderedWordCount failed", job.run(tezConf, new String[]{
            inputDirStr, outputDirStr, "5"}, null) == 0);
  }
  TestTezJobs.verifyOutput(outputDir, remoteFs);
  List<HistoryEvent> historyEventsOfAttempt1 = RecoveryParser
      .readRecoveryEvents(tezConf, job.getAppId(), 1);
  HistoryEvent lastEvent = historyEventsOfAttempt1
      .get(historyEventsOfAttempt1.size() - 1);
  assertEquals(shutdownCondition.getEvent().getEventType(),
      lastEvent.getEventType());
  assertTrue(shutdownCondition.match(lastEvent));

}
 

@Test(timeout = 100000)
@SuppressWarnings("unchecked")
public void testSubmitDagInSessionWithLargeDagPlan() throws Exception {
  int maxIPCMsgSize = 1024;
  String dagPlanName = "dagplan-name";
  File requestFile = tmpFolder.newFile("request-file");
  TezConfiguration conf = new TezConfiguration();
  conf.setInt(CommonConfigurationKeys.IPC_MAXIMUM_DATA_LENGTH, maxIPCMsgSize);

  // Check with 70 MB (64 MB is CodedInputStream's default limit in earlier versions of protobuf)
  byte[] randomBytes = new byte[70 << 20];
  (new Random()).nextBytes(randomBytes);
  UserPayload payload = UserPayload.create(ByteBuffer.wrap(randomBytes));
  Vertex vertex = Vertex.create("V", ProcessorDescriptor.create("P").setUserPayload(payload), 1);
  DAGPlan dagPlan = DAG.create(dagPlanName).addVertex(vertex).createDag(conf, null, null, null, false);

  String lrName = "localResource";
  String scheme = "file";
  String host = "localhost";
  int port = 80;
  String path = "/test";
  URL lrURL = URL.newInstance(scheme, host, port, path);
  LocalResource localResource = LocalResource.newInstance(lrURL, LocalResourceType.FILE,
      LocalResourceVisibility.PUBLIC, 1, 1);
  Map<String, LocalResource> localResources = new HashMap<>();
  localResources.put(lrName, localResource);

  SubmitDAGRequestProto.Builder requestBuilder = SubmitDAGRequestProto.newBuilder().setDAGPlan(dagPlan)
      .setAdditionalAmResources(DagTypeConverters.convertFromLocalResources(localResources));
  try (FileOutputStream fileOutputStream = new FileOutputStream(requestFile)) {
    requestBuilder.build().writeTo(fileOutputStream);
  }

  DAGClientHandler dagClientHandler = mock(DAGClientHandler.class);
  ACLManager aclManager = mock(ACLManager.class);
  DAGClientAMProtocolBlockingPBServerImpl serverImpl = spy(new DAGClientAMProtocolBlockingPBServerImpl(
      dagClientHandler, FileSystem.get(conf)));
  when(dagClientHandler.getACLManager()).thenReturn(aclManager);
  when(dagClientHandler.submitDAG((DAGPlan)any(), (Map<String, LocalResource>)any())).thenReturn("dag-id");
  when(aclManager.checkAMModifyAccess((UserGroupInformation) any())).thenReturn(true);

  requestBuilder.clear().setSerializedRequestPath(requestFile.getAbsolutePath());
  serverImpl.submitDAG(null, requestBuilder.build());

  ArgumentCaptor<DAGPlan> dagPlanCaptor = ArgumentCaptor.forClass(DAGPlan.class);
  verify(dagClientHandler).submitDAG(dagPlanCaptor.capture(), localResourcesCaptor.capture());
  dagPlan = dagPlanCaptor.getValue();
  localResources = localResourcesCaptor.getValue();

  assertEquals(dagPlan.getName(), dagPlanName);
  assertEquals(dagPlan.getVertexCount(), 1);
  assertTrue(dagPlan.getSerializedSize() > maxIPCMsgSize);
  assertArrayEquals(randomBytes, dagPlan.getVertex(0).getProcessorDescriptor().getTezUserPayload().getUserPayload().
      toByteArray());
  assertEquals(localResources.size(), 1);
  assertTrue(localResources.containsKey(lrName));
  localResource = localResources.get(lrName);
  assertEquals(localResource.getType(), LocalResourceType.FILE);
  assertEquals(localResource.getVisibility(), LocalResourceVisibility.PUBLIC);
  lrURL = localResource.getResource();
  assertEquals(lrURL.getScheme(), scheme);
  assertEquals(lrURL.getHost(), host);
  assertEquals(lrURL.getPort(), port);
  assertEquals(lrURL.getFile(), path);
}
 

private static void adjustAMConfig(TezConfiguration amConf, TezJobConfig tezJobConf) {
    int requiredAMMaxHeap = -1;
    int requiredAMResourceMB = -1;
    int configuredAMMaxHeap = Utils.extractHeapSizeInMB(amConf.get(
            TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS,
            TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS_DEFAULT));
    int configuredAMResourceMB = amConf.getInt(
            TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB,
            TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB_DEFAULT);

    if (tezJobConf.getEstimatedTotalParallelism() > 0) {

        int minAMMaxHeap = 3584;
        int minAMResourceMB = 4096;

        // Rough estimation. For 5K tasks 1G Xmx and 1.5G resource.mb
        // Increment by 512 mb for every additional 5K tasks.
        for (int taskCount = 30000; taskCount >= 5000; taskCount-=5000) {
            if (tezJobConf.getEstimatedTotalParallelism() > taskCount) {
                requiredAMMaxHeap = minAMMaxHeap;
                requiredAMResourceMB = minAMResourceMB;
                break;
            }
            minAMMaxHeap = minAMMaxHeap - 512;
            minAMResourceMB = minAMResourceMB - 512;
        }

        if (requiredAMResourceMB > -1 && configuredAMResourceMB < requiredAMResourceMB) {
            amConf.setInt(TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, requiredAMResourceMB);
            log.info("Increasing "
                    + TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB + " from "
                    + configuredAMResourceMB + " to "
                    + requiredAMResourceMB
                    + " as the number of total estimated tasks is "
                    + tezJobConf.getEstimatedTotalParallelism());

            if (requiredAMMaxHeap > -1 && configuredAMMaxHeap < requiredAMMaxHeap) {
                amConf.set(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS,
                        amConf.get(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS)
                                + " -Xmx" + requiredAMMaxHeap + "M");
                log.info("Increasing Tez AM Heap Size from "
                        + configuredAMMaxHeap + "M to "
                        + requiredAMMaxHeap
                        + "M as the number of total estimated tasks is "
                        + tezJobConf.getEstimatedTotalParallelism());
            }
        }
    }
}