xref: /aosp_15_r20/external/pytorch/tools/testing/target_determination/heuristics/profiling.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1from __future__ import annotations
2
3from typing import Any
4
5from tools.stats.import_test_stats import (
6    ADDITIONAL_CI_FILES_FOLDER,
7    TD_HEURISTIC_PROFILING_FILE,
8)
9from tools.testing.target_determination.heuristics.interface import (
10    HeuristicInterface,
11    TestPrioritizations,
12)
13from tools.testing.target_determination.heuristics.utils import (
14    get_ratings_for_tests,
15    normalize_ratings,
16)
17from tools.testing.test_run import TestRun
18
19
20# Profilers were used to gather simple python code coverage information for each
21# test to see files were involved in each tests and used to build a correlation
22# dict (where all ratings are 1).
23class Profiling(HeuristicInterface):
24    def __init__(self, **kwargs: Any) -> None:
25        super().__init__(**kwargs)
26
27    def get_prediction_confidence(self, tests: list[str]) -> TestPrioritizations:
28        test_ratings = get_ratings_for_tests(
29            ADDITIONAL_CI_FILES_FOLDER / TD_HEURISTIC_PROFILING_FILE
30        )
31        test_ratings = {TestRun(k): v for (k, v) in test_ratings.items() if k in tests}
32        return TestPrioritizations(tests, normalize_ratings(test_ratings, 0.25))
33