bazel_adb_run.py

#!/usr/bin/env python

# Must run at root dir of libmace project.
# python tools/bazel_adb_run.py \
#     --target_abis=armeabi-v7a \
#     --target_socs=sdm845
#     --target=//mace/ops:ops_test
#     --stdout_processor=stdout_processor

import argparse
import random
import re
import sys

import sh_commands


def stdout_processor(stdout, device_properties, abi):
    pass


def unittest_stdout_processor(stdout, device_properties, abi):
    stdout_lines = stdout.split("\n")
    for line in stdout_lines:
        if "Aborted" in line or "FAILED" in line:
            raise Exception("Command failed")


def ops_benchmark_stdout_processor(stdout, device_properties, abi):
    stdout_lines = stdout.split("\n")
    metrics = {}
    for line in stdout_lines:
        if "Aborted" in line:
            raise Exception("Command failed")
        line = line.strip()
        parts = line.split()
        if len(parts) == 5 and parts[0].startswith("BM_"):
            metrics["%s.time_ms" % parts[0]] = str(float(parts[1]) / 1e6)
            metrics["%s.input_mb_per_sec" % parts[0]] = parts[3]
            metrics["%s.gmacc_per_sec" % parts[0]] = parts[4]

    platform = device_properties["ro.board.platform"].replace(" ", "-")
    model = device_properties["ro.product.model"].replace(" ", "-")
    tags = {
        "ro.board.platform": platform,
        "ro.product.model": model,
        "abi": abi
    }
    sh_commands.falcon_push_metrics(
        metrics, tags=tags, endpoint="mace_ops_benchmark")


def parse_args():
    """Parses command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--target_abis",
        type=str,
        default="armeabi-v7a",
        help="Target ABIs, comma seperated list")
    parser.add_argument(
        "--target_socs",
        type=str,
        default="all",
        help="SoCs (ro.board.platform from getprop) to build, "
        "comma seperated list or all/random")
    parser.add_argument(
        "--target", type=str, default="//...", help="Bazel target to build")
    parser.add_argument(
        "--run_target",
        type=bool,
        default=False,
        help="Whether to run the target")
    parser.add_argument("--args", type=str, default="", help="Command args")
    parser.add_argument(
        "--stdout_processor",
        type=str,
        default="stdout_processor",
        help="Stdout processing function, default: stdout_processor")
    return parser.parse_known_args()


def main(unused_args):
    target_socs = None
    if FLAGS.target_socs != "all" and FLAGS.target_socs != "random":
        target_socs = set(FLAGS.target_socs.split(','))
    target_devices = sh_commands.adb_devices(target_socs=target_socs)
    if FLAGS.target_socs == "random":
        target_devices = [random.choice(target_devices)]

    target = FLAGS.target
    host_bin_path, bin_name = sh_commands.bazel_target_to_bin(target)
    target_abis = FLAGS.target_abis.split(',')

    # generate sources
    sh_commands.gen_encrypted_opencl_source()
    sh_commands.gen_compiled_opencl_source()
    sh_commands.gen_mace_version()

    for target_abi in target_abis:
        sh_commands.bazel_build(target, abi=target_abi)
        if FLAGS.run_target:
            for serialno in target_devices:
                if target_abi not in set(
                        sh_commands.adb_supported_abis(serialno)):
                    print("Skip device %s which does not support ABI %s" %
                          (serialno, target_abi))
                    continue
                stdouts = sh_commands.adb_run(
                    serialno,
                    host_bin_path,
                    bin_name,
                    args=FLAGS.args,
                    opencl_profiling=1,
                    vlog_level=0,
                    device_bin_path="/data/local/tmp/mace",
                    out_of_range_check=1)
                device_properties = sh_commands.adb_getprop_by_serialno(
                    serialno)
                globals()[FLAGS.stdout_processor](stdouts, device_properties,
                                                  target_abi)


if __name__ == "__main__":
    FLAGS, unparsed = parse_args()
    main(unused_args=[sys.argv[0]] + unparsed)