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