1from __future__ import annotations 2 3import os 4import subprocess 5import time 6 7from ..util.setting import ( 8 JSON_FOLDER_BASE_DIR, 9 MERGED_FOLDER_BASE_DIR, 10 TestList, 11 TestPlatform, 12 TestType, 13) 14from ..util.utils import ( 15 check_platform_type, 16 convert_to_relative_path, 17 create_folder, 18 get_raw_profiles_folder, 19 get_test_name_from_whole_path, 20 print_log, 21 print_time, 22 related_to_test_list, 23 replace_extension, 24) 25from .utils import get_tool_path_by_platform, run_cpp_test 26 27 28def create_corresponding_folder( 29 cur_path: str, prefix_cur_path: str, dir_list: list[str], new_base_folder: str 30) -> None: 31 for dir_name in dir_list: 32 relative_path = convert_to_relative_path( 33 cur_path, prefix_cur_path 34 ) # get folder name like 'aten' 35 new_folder_path = os.path.join(new_base_folder, relative_path, dir_name) 36 create_folder(new_folder_path) 37 38 39def run_target( 40 binary_file: str, raw_file: str, test_type: TestType, platform_type: TestPlatform 41) -> None: 42 print_log("start run: ", binary_file) 43 # set environment variable -- raw profile output path of the binary run 44 os.environ["LLVM_PROFILE_FILE"] = raw_file 45 # run binary 46 if test_type == TestType.PY and platform_type == TestPlatform.OSS: 47 from ..oss.utils import run_oss_python_test 48 49 run_oss_python_test(binary_file) 50 else: 51 run_cpp_test(binary_file) 52 53 54def merge_target(raw_file: str, merged_file: str, platform_type: TestPlatform) -> None: 55 print_log("start to merge target: ", raw_file) 56 # run command 57 llvm_tool_path = get_tool_path_by_platform(platform_type) 58 subprocess.check_call( 59 [ 60 f"{llvm_tool_path}/llvm-profdata", 61 "merge", 62 "-sparse", 63 raw_file, 64 "-o", 65 merged_file, 66 ] 67 ) 68 69 70def export_target( 71 merged_file: str, 72 json_file: str, 73 binary_file: str, 74 shared_library_list: list[str], 75 platform_type: TestPlatform, 76) -> None: 77 if binary_file is None: 78 raise Exception( # noqa: TRY002 79 f"{merged_file} doesn't have corresponding binary!" 80 ) # noqa: TRY002 81 print_log("start to export: ", merged_file) 82 # run export 83 cmd_shared_library = ( 84 "" 85 if not shared_library_list 86 else f" -object {' -object '.join(shared_library_list)}" 87 ) 88 # if binary_file = "", then no need to add it (python test) 89 cmd_binary = "" if not binary_file else f" -object {binary_file} " 90 llvm_tool_path = get_tool_path_by_platform(platform_type) 91 92 cmd = f"{llvm_tool_path}/llvm-cov export {cmd_binary} {cmd_shared_library} -instr-profile={merged_file} > {json_file}" 93 os.system(cmd) 94 95 96def merge(test_list: TestList, platform_type: TestPlatform) -> None: 97 print("start merge") 98 start_time = time.time() 99 # find all raw profile under raw_folder and sub-folders 100 raw_folder_path = get_raw_profiles_folder() 101 g = os.walk(raw_folder_path) 102 for path, dir_list, file_list in g: 103 # if there is a folder raw/aten/, create corresponding merged folder profile/merged/aten/ if not exists yet 104 create_corresponding_folder( 105 path, raw_folder_path, dir_list, MERGED_FOLDER_BASE_DIR 106 ) 107 # check if we can find raw profile under this path's folder 108 for file_name in file_list: 109 if file_name.endswith(".profraw"): 110 if not related_to_test_list(file_name, test_list): 111 continue 112 print(f"start merge {file_name}") 113 raw_file = os.path.join(path, file_name) 114 merged_file_name = replace_extension(file_name, ".merged") 115 merged_file = os.path.join( 116 MERGED_FOLDER_BASE_DIR, 117 convert_to_relative_path(path, raw_folder_path), 118 merged_file_name, 119 ) 120 merge_target(raw_file, merged_file, platform_type) 121 print_time("merge take time: ", start_time, summary_time=True) 122 123 124def export(test_list: TestList, platform_type: TestPlatform) -> None: 125 print("start export") 126 start_time = time.time() 127 # find all merged profile under merged_folder and sub-folders 128 g = os.walk(MERGED_FOLDER_BASE_DIR) 129 for path, dir_list, file_list in g: 130 # create corresponding merged folder in [json folder] if not exists yet 131 create_corresponding_folder( 132 path, MERGED_FOLDER_BASE_DIR, dir_list, JSON_FOLDER_BASE_DIR 133 ) 134 # check if we can find merged profile under this path's folder 135 for file_name in file_list: 136 if file_name.endswith(".merged"): 137 if not related_to_test_list(file_name, test_list): 138 continue 139 print(f"start export {file_name}") 140 # merged file 141 merged_file = os.path.join(path, file_name) 142 # json file 143 json_file_name = replace_extension(file_name, ".json") 144 json_file = os.path.join( 145 JSON_FOLDER_BASE_DIR, 146 convert_to_relative_path(path, MERGED_FOLDER_BASE_DIR), 147 json_file_name, 148 ) 149 check_platform_type(platform_type) 150 # binary file and shared library 151 binary_file = "" 152 shared_library_list = [] 153 if platform_type == TestPlatform.FBCODE: 154 from caffe2.fb.code_coverage.tool.package.fbcode.utils import ( # type: ignore[import] 155 get_fbcode_binary_folder, 156 ) 157 158 binary_file = os.path.join( 159 get_fbcode_binary_folder(path), 160 get_test_name_from_whole_path(merged_file), 161 ) 162 elif platform_type == TestPlatform.OSS: 163 from ..oss.utils import get_oss_binary_file, get_oss_shared_library 164 165 test_name = get_test_name_from_whole_path(merged_file) 166 # if it is python test, no need to provide binary, shared library is enough 167 binary_file = ( 168 "" 169 if test_name.endswith(".py") 170 else get_oss_binary_file(test_name, TestType.CPP) 171 ) 172 shared_library_list = get_oss_shared_library() 173 export_target( 174 merged_file, 175 json_file, 176 binary_file, 177 shared_library_list, 178 platform_type, 179 ) 180 print_time("export take time: ", start_time, summary_time=True) 181