utils/benchmark/tools/gbench/util.py

   1 """util.py - General utilities for running, loading, and processing benchmarks
   2 """
   3 import json
   4 import os
   5 import tempfile
   6 import subprocess
   7 import sys
   8
   9 # Input file type enumeration
  10 IT_Invalid    = 0
  11 IT_JSON       = 1
  12 IT_Executable = 2
  13
  14 _num_magic_bytes = 2 if sys.platform.startswith('win') else 4
  15 def is_executable_file(filename):
  16     """
  17     Return 'True' if 'filename' names a valid file which is likely
  18     an executable. A file is considered an executable if it starts with the
  19     magic bytes for a EXE, Mach O, or ELF file.
  20     """
  21     if not os.path.isfile(filename):
  22         return False
  23     with open(filename, mode='rb') as f:
  24         magic_bytes = f.read(_num_magic_bytes)
  25     if sys.platform == 'darwin':
  26         return magic_bytes in [
  27             b'\xfe\xed\xfa\xce',  # MH_MAGIC
  28             b'\xce\xfa\xed\xfe',  # MH_CIGAM
  29             b'\xfe\xed\xfa\xcf',  # MH_MAGIC_64
  30             b'\xcf\xfa\xed\xfe',  # MH_CIGAM_64
  31             b'\xca\xfe\xba\xbe',  # FAT_MAGIC
  32             b'\xbe\xba\xfe\xca'   # FAT_CIGAM
  33         ]
  34     elif sys.platform.startswith('win'):
  35         return magic_bytes == b'MZ'
  36     else:
  37         return magic_bytes == b'\x7FELF'
  38
  39
  40 def is_json_file(filename):
  41     """
  42     Returns 'True' if 'filename' names a valid JSON output file.
  43     'False' otherwise.
  44     """
  45     try:
  46         with open(filename, 'r') as f:
  47             json.load(f)
  48         return True
  49     except:
  50         pass
  51     return False
  52
  53
  54 def classify_input_file(filename):
  55     """
  56     Return a tuple (type, msg) where 'type' specifies the classified type
  57     of 'filename'. If 'type' is 'IT_Invalid' then 'msg' is a human readable
  58     string represeting the error.
  59     """
  60     ftype = IT_Invalid
  61     err_msg = None
  62     if not os.path.exists(filename):
  63         err_msg = "'%s' does not exist" % filename
  64     elif not os.path.isfile(filename):
  65         err_msg = "'%s' does not name a file" % filename
  66     elif is_executable_file(filename):
  67         ftype = IT_Executable
  68     elif is_json_file(filename):
  69         ftype = IT_JSON
  70     else:
  71         err_msg = "'%s' does not name a valid benchmark executable or JSON file" % filename
  72     return ftype, err_msg
  73
  74
  75 def check_input_file(filename):
  76     """
  77     Classify the file named by 'filename' and return the classification.
  78     If the file is classified as 'IT_Invalid' print an error message and exit
  79     the program.
  80     """
  81     ftype, msg = classify_input_file(filename)
  82     if ftype == IT_Invalid:
  83         print("Invalid input file: %s" % msg)
  84         sys.exit(1)
  85     return ftype
  86
  87 def find_benchmark_flag(prefix, benchmark_flags):
  88     """
  89     Search the specified list of flags for a flag matching `<prefix><arg>` and
  90     if it is found return the arg it specifies. If specified more than once the
  91     last value is returned. If the flag is not found None is returned.
  92     """
  93     assert prefix.startswith('--') and prefix.endswith('=')
  94     result = None
  95     for f in benchmark_flags:
  96         if f.startswith(prefix):
  97             result = f[len(prefix):]
  98     return result
  99
 100 def remove_benchmark_flags(prefix, benchmark_flags):
 101     """
 102     Return a new list containing the specified benchmark_flags except those
 103     with the specified prefix.
 104     """
 105     assert prefix.startswith('--') and prefix.endswith('=')
 106     return [f for f in benchmark_flags if not f.startswith(prefix)]
 107
 108 def load_benchmark_results(fname):
 109     """
 110     Read benchmark output from a file and return the JSON object.
 111     REQUIRES: 'fname' names a file containing JSON benchmark output.
 112     """
 113     with open(fname, 'r') as f:
 114         return json.load(f)
 115
 116
 117 def run_benchmark(exe_name, benchmark_flags):
 118     """
 119     Run a benchmark specified by 'exe_name' with the specified
 120     'benchmark_flags'. The benchmark is run directly as a subprocess to preserve
 121     real time console output.
 122     RETURNS: A JSON object representing the benchmark output
 123     """
 124     output_name = find_benchmark_flag('--benchmark_out=',
 125                                       benchmark_flags)
 126     is_temp_output = False
 127     if output_name is None:
 128         is_temp_output = True
 129         thandle, output_name = tempfile.mkstemp()
 130         os.close(thandle)
 131         benchmark_flags = list(benchmark_flags) + \
 132                           ['--benchmark_out=%s' % output_name]
 133
 134     cmd = [exe_name] + benchmark_flags
 135     print("RUNNING: %s" % ' '.join(cmd))
 136     exitCode = subprocess.call(cmd)
 137     if exitCode != 0:
 138         print('TEST FAILED...')
 139         sys.exit(exitCode)
 140     json_res = load_benchmark_results(output_name)
 141     if is_temp_output:
 142         os.unlink(output_name)
 143     return json_res
 144
 145
 146 def run_or_load_benchmark(filename, benchmark_flags):
 147     """
 148     Get the results for a specified benchmark. If 'filename' specifies
 149     an executable benchmark then the results are generated by running the
 150     benchmark. Otherwise 'filename' must name a valid JSON output file,
 151     which is loaded and the result returned.
 152     """
 153     ftype = check_input_file(filename)
 154     if ftype == IT_JSON:
 155         return load_benchmark_results(filename)
 156     elif ftype == IT_Executable:
 157         return run_benchmark(filename, benchmark_flags)
 158     else:
 159         assert False # This branch is unreachable