org.apache.tez.test.FaultToleranceTestRunner.java Source code

Introduction

Here is the source code for org.apache.tez.test.FaultToleranceTestRunner.java

Source

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.tez.test;

import java.lang.reflect.Method;
import java.util.Random;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.tez.client.TezClientUtils;
import org.apache.tez.client.TezClient;
import org.apache.tez.dag.api.DAG;
import org.apache.tez.dag.api.TezConfiguration;
import org.apache.tez.dag.api.client.DAGClient;
import org.apache.tez.dag.api.client.DAGStatus;
import org.apache.hadoop.util.GenericOptionsParser;

/**
 * Run a DAG on a cluster with the given configuration. Starts a TezSession
 * using default cluster configuration from installation. Then uses reflection
 * to get the class from the first class-name argument. That class must have a
 * static method - createDAG(org.apache.hadoop.conf.Configuration) that returns
 * a DAG. Configuration is picked up by reading the file specified via the
 * second path argument. The static method is invoked to get the DAG. The DAG is
 * then executed in the session. Returns success if DAG succeeds.
 */
public class FaultToleranceTestRunner {

    static String TEST_ROOT_DIR = "tmp";
    Configuration conf = null;
    TezClient tezSession = null;
    Resource defaultResource = Resource.newInstance(100, 0);

    void setup() throws Exception {
        TezConfiguration tezConf = null;
        if (conf == null) {
            tezConf = new TezConfiguration(new YarnConfiguration());
        } else {
            tezConf = new TezConfiguration(new YarnConfiguration(this.conf));
        }
        FileSystem defaultFs = FileSystem.get(tezConf);

        Path remoteStagingDir = defaultFs
                .makeQualified(new Path(TEST_ROOT_DIR, String.valueOf(new Random().nextInt(100000))));
        TezClientUtils.ensureStagingDirExists(tezConf, remoteStagingDir);

        tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, remoteStagingDir.toString());

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

    void tearDown() throws Exception {
        if (tezSession != null) {
            tezSession.stop();
        }
    }

    DAG getDAG(String className, String confFilePath) throws Exception {
        Class<?> clazz = Class.forName(className);
        Method method = clazz.getMethod("createDAG", Configuration.class);

        Configuration testConf = new Configuration(false);
        if (confFilePath != null) {
            Path confPath = new Path(confFilePath);
            testConf.addResource(confPath);
        }

        DAG dag = (DAG) method.invoke(null, testConf);

        return dag;
    }

    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 appContext: " + dagClient.getExecutionContext() + " Current state: "
                        + dagStatus.getState());
                Thread.sleep(500);
                dagStatus = dagClient.getDAGStatus(null);
            }

            if (dagStatus.getState() == DAGStatus.State.SUCCEEDED) {
                return true;
            }

        } finally {
            tearDown();
        }

        return false;
    }

    static void printUsage() {
        System.err.println(
                "Usage: " + " FaultToleranceTestRunner [generic options] <dag-class-name> <test-conf-path>");
        GenericOptionsParser.printGenericCommandUsage(System.err);
    }

    public static void main(String[] args) throws Exception {
        Configuration conf = new Configuration();
        String[] otherArgs = new GenericOptionsParser(conf, args).getRemainingArgs();
        String className = null;
        String confFilePath = null;
        if (otherArgs.length == 1) {
            className = otherArgs[0];
        } else if (otherArgs.length == 2) {
            className = otherArgs[0];
            confFilePath = otherArgs[1];
        } else {
            printUsage();
            System.exit(1);
        }

        FaultToleranceTestRunner job = new FaultToleranceTestRunner();
        if (job.run(conf, className, confFilePath)) {
            System.out.println("Succeeded.");
        } else {
            System.out.println("Failed.");
            System.exit(2);
        }
    }
}