xref: /aosp_15_r20/external/pytorch/tools/code_coverage/package/oss/init.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1from __future__ import annotations
2
3import argparse
4import os
5from typing import cast
6
7from ..util.setting import (
8    CompilerType,
9    JSON_FOLDER_BASE_DIR,
10    LOG_DIR,
11    Option,
12    Test,
13    TestList,
14    TestType,
15)
16from ..util.utils import (
17    clean_up,
18    create_folder,
19    print_log,
20    raise_no_test_found_exception,
21    remove_file,
22    remove_folder,
23)
24from ..util.utils_init import add_arguments_utils, create_folders, get_options
25from .utils import (
26    clean_up_gcda,
27    detect_compiler_type,
28    get_llvm_tool_path,
29    get_oss_binary_folder,
30    get_pytorch_folder,
31)
32
33
34BLOCKED_PYTHON_TESTS = {
35    "run_test.py",
36    "test_dataloader.py",
37    "test_multiprocessing.py",
38    "test_multiprocessing_spawn.py",
39    "test_utils.py",
40}
41
42
43def initialization() -> tuple[Option, TestList, list[str]]:
44    # create folder if not exists
45    create_folders()
46    # add arguments
47    parser = argparse.ArgumentParser()
48    parser = add_arguments_utils(parser)
49    parser = add_arguments_oss(parser)
50    # parse arguments
51    (options, args_interested_folder, args_run_only, arg_clean) = parse_arguments(
52        parser
53    )
54    # clean up
55    if arg_clean:
56        clean_up_gcda()
57        clean_up()
58    # get test lists
59    test_list = get_test_list(args_run_only)
60    # get interested folder -- final report will only over these folders
61    interested_folders = empty_list_if_none(args_interested_folder)
62    # print initialization information
63    print_init_info()
64    # remove last time's log
65    remove_file(os.path.join(LOG_DIR, "log.txt"))
66    return (options, test_list, interested_folders)
67
68
69def add_arguments_oss(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
70    parser.add_argument(
71        "--run-only",
72        help="only run certain test(s), for example: atest test_nn.py.",
73        nargs="*",
74        default=None,
75    )
76
77    return parser
78
79
80def parse_arguments(
81    parser: argparse.ArgumentParser,
82) -> tuple[Option, list[str] | None, list[str] | None, bool | None]:
83    # parse args
84    args = parser.parse_args()
85    # get option
86    options = get_options(args)
87    return (options, args.interest_only, args.run_only, args.clean)
88
89
90def get_test_list_by_type(run_only: list[str] | None, test_type: TestType) -> TestList:
91    test_list: TestList = []
92    binary_folder = get_oss_binary_folder(test_type)
93    g = os.walk(binary_folder)
94    for _, _, file_list in g:
95        for file_name in file_list:
96            if run_only is not None and file_name not in run_only:
97                continue
98            # target pattern in oss is used in printing report -- which tests we have run
99            test: Test = Test(
100                name=file_name,
101                target_pattern=file_name,
102                test_set="",
103                test_type=test_type,
104            )
105            test_list.append(test)
106    return test_list
107
108
109def get_test_list(run_only: list[str] | None) -> TestList:
110    test_list: TestList = []
111    # add c++ test list
112    test_list.extend(get_test_list_by_type(run_only, TestType.CPP))
113    # add python test list
114    py_run_only = get_python_run_only(run_only)
115    test_list.extend(get_test_list_by_type(py_run_only, TestType.PY))
116
117    # not find any test to run
118    if not test_list:
119        raise_no_test_found_exception(
120            get_oss_binary_folder(TestType.CPP), get_oss_binary_folder(TestType.PY)
121        )
122    return test_list
123
124
125def empty_list_if_none(arg_interested_folder: list[str] | None) -> list[str]:
126    if arg_interested_folder is None:
127        return []
128    # if this argument is specified, just return itself
129    return arg_interested_folder
130
131
132def gcc_export_init() -> None:
133    remove_folder(JSON_FOLDER_BASE_DIR)
134    create_folder(JSON_FOLDER_BASE_DIR)
135
136
137def get_python_run_only(args_run_only: list[str] | None) -> list[str]:
138    # if user specifies run-only option
139    if args_run_only:
140        return args_run_only
141
142    # if not specified, use default setting, different for gcc and clang
143    if detect_compiler_type() == CompilerType.GCC:
144        return ["run_test.py"]
145    else:
146        # for clang, some tests will result in too large intermediate files that can't be merged by llvm, we need to skip them
147        run_only: list[str] = []
148        binary_folder = get_oss_binary_folder(TestType.PY)
149        g = os.walk(binary_folder)
150        for _, _, file_list in g:
151            for file_name in file_list:
152                if file_name in BLOCKED_PYTHON_TESTS or not file_name.endswith(".py"):
153                    continue
154                run_only.append(file_name)
155            # only run tests in the first-level folder in test/
156            break
157        return run_only
158
159
160def print_init_info() -> None:
161    print_log("pytorch folder: ", get_pytorch_folder())
162    print_log("cpp test binaries folder: ", get_oss_binary_folder(TestType.CPP))
163    print_log("python test scripts folder: ", get_oss_binary_folder(TestType.PY))
164    print_log("compiler type: ", cast(CompilerType, detect_compiler_type()).value)
165    print_log(
166        "llvm tool folder (only for clang, if you are using gcov please ignore it): ",
167        get_llvm_tool_path(),
168    )
169