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
