pytorch/tools/testing/test_selections.py

import csv
import json
import os
import subprocess

from tools.stats.s3_stat_parser import (
    get_previous_reports_for_branch,
    get_previous_reports_for_pr,
    Report, Version2Report,
    HAVE_BOTO3)
from tools.stats.import_test_stats import (
    get_disabled_tests,
    get_slow_tests
)

from typing import Any, Dict, List, Optional, Tuple, cast
from typing_extensions import TypedDict

class JobTimeJSON(TypedDict):
    commit: str
    JOB_BASE_NAME: str
    job_times: Dict[str, float]


def _get_stripped_CI_job() -> str:
    """E.g. convert 'pytorch_windows_vs2019_py36_cuda10.1_build' to 'pytorch_windows_vs2019_py36_cuda10.1'.
    """
    job = os.environ.get("JOB_BASE_NAME", "").rstrip('0123456789')
    if job.endswith('_slow_test'):
        job = job[:len(job) - len('_slow_test')]
    elif job.endswith('_test') or job.endswith('-test'):
        job = job[:len(job) - len('_test')]
    elif job.endswith('_build') or job.endswith('-build'):
        job = job[:len(job) - len('_build')]
    return job


def _get_job_times_json(job_times: Dict[str, float]) -> JobTimeJSON:
    return {
        'commit': subprocess.check_output(['git', 'rev-parse', 'HEAD'], encoding="ascii").strip(),
        'JOB_BASE_NAME': _get_stripped_CI_job(),
        'job_times': job_times,
    }


def _calculate_job_times(reports: List["Report"]) -> Dict[str, float]:
    """Compute test runtime by filename: ("test_file_name" -> (current_avg, # values))
    """
    jobs_to_times: Dict[str, Tuple[float, int]] = dict()
    for report in reports:
        v_report = cast(Version2Report, report)
        assert 'format_version' in v_report.keys() and v_report.get('format_version') == 2, \
            "S3 format currently handled is version 2 only"
        files: Dict[str, Any] = v_report['files']
        for name, test_file in files.items():
            if name not in jobs_to_times:
                jobs_to_times[name] = (test_file['total_seconds'], 1)
            else:
                curr_avg, curr_count = jobs_to_times[name]
                new_count = curr_count + 1
                new_avg = (curr_avg * curr_count + test_file['total_seconds']) / new_count
                jobs_to_times[name] = (new_avg, new_count)

    return {job: time for job, (time, _) in jobs_to_times.items()}


def calculate_shards(num_shards: int, tests: List[str], job_times: Dict[str, float]) -> List[Tuple[float, List[str]]]:
    filtered_job_times: Dict[str, float] = dict()
    unknown_jobs : List[str] = []
    for test in tests:
        if test in job_times:
            filtered_job_times[test] = job_times[test]
        else:
            unknown_jobs.append(test)

    # The following attempts to implement a partition approximation greedy algorithm
    # See more at https://en.wikipedia.org/wiki/Greedy_number_partitioning
    sorted_jobs = sorted(filtered_job_times, key=lambda j: filtered_job_times[j], reverse=True)
    sharded_jobs: List[Tuple[float, List[str]]] = [(0.0, []) for _ in range(num_shards)]
    for job in sorted_jobs:
        min_shard_index = sorted(range(num_shards), key=lambda i: sharded_jobs[i][0])[0]
        curr_shard_time, curr_shard_jobs = sharded_jobs[min_shard_index]
        curr_shard_jobs.append(job)
        sharded_jobs[min_shard_index] = (curr_shard_time + filtered_job_times[job], curr_shard_jobs)

    # Round robin the unknown jobs starting with the smallest shard
    index = sorted(range(num_shards), key=lambda i: sharded_jobs[i][0])[0]
    for job in unknown_jobs:
        sharded_jobs[index][1].append(job)
        index = (index + 1) % num_shards
    return sharded_jobs


def _pull_job_times_from_S3() -> Dict[str, float]:
    if HAVE_BOTO3:
        ci_job_prefix = _get_stripped_CI_job()
        s3_reports: List["Report"] = get_previous_reports_for_branch('origin/viable/strict', ci_job_prefix)
    else:
        print('Uh oh, boto3 is not found. Either it is not installed or we failed to import s3_stat_parser.')
        print('If not installed, please install boto3 for automatic sharding and test categorization.')
        s3_reports = []

    if len(s3_reports) == 0:
        print('Gathered no reports from S3. Please proceed without them.')
        return dict()

    return _calculate_job_times(s3_reports)


def _query_past_job_times(test_times_file: Optional[str] = None) -> Dict[str, float]:
    """Read historic test job times from a file.

    If the file doesn't exist or isn't matching current commit. It will download data from S3 and exported it.
    """
    if test_times_file and os.path.exists(test_times_file):
        with open(test_times_file) as file:
            test_times_json: JobTimeJSON = json.load(file)

        curr_commit = subprocess.check_output(['git', 'rev-parse', 'HEAD'], encoding="ascii").strip()
        file_commit = test_times_json.get('commit', '')
        curr_ci_job = _get_stripped_CI_job()
        file_ci_job = test_times_json.get('JOB_BASE_NAME', 'N/A')
        if curr_commit != file_commit:
            print(f'Current test times file is from different commit {file_commit}.')
        elif curr_ci_job != file_ci_job:
            print(f'Current test times file is for different CI job {file_ci_job}.')
        else:
            print(f'Found stats for current commit: {curr_commit} and job: {curr_ci_job}. Proceeding with those values.')
            return test_times_json.get('job_times', {})

        # Found file, but commit or CI job in JSON doesn't match
        print(f'Overwriting current file with stats based on current commit: {curr_commit} and CI job: {curr_ci_job}')

    job_times = export_S3_test_times(test_times_file)

    return job_times


def _query_failure_test_module(reports: List[Tuple["Report", str]]) -> List[str]:
    test_modules: List[str] = []
    if len(reports) == 0 or len(reports[0]) == 0:
        return test_modules
    report = reports[0][0]
    v_report = cast(Version2Report, report)
    assert 'format_version' in v_report.keys() and v_report.get('format_version') == 2, \
        "S3 format currently handled is version 2 only"
    files: Dict[str, Any] = v_report['files']
    for fname, file in files.items():
        contains_failure = any(
            any(case['status'] == 'errored' or case['status'] == 'failed'
                for _, case in suite['cases'].items())
            for _, suite in file['suites'].items())
        if contains_failure:
            test_modules.append(fname)
    return test_modules


def _query_changed_test_files() -> List[str]:
    cmd = ["git", "diff", "--name-only", "origin/master", "HEAD"]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    if proc.returncode != 0:
        raise RuntimeError("Unable to get changed files")

    lines = proc.stdout.decode().strip().split("\n")
    lines = [line.strip() for line in lines]
    return lines


def get_shard_based_on_S3(which_shard: int, num_shards: int, tests: List[str], test_times_file: str) -> List[str]:
    """Get sharded test allocation based on historic S3 data.
    """
    jobs_to_times = _query_past_job_times(test_times_file)

    # Got no stats from S3, returning early to save runtime
    if len(jobs_to_times) == 0:
        print('Gathered no stats from S3. Proceeding with default sharding plan.')
        return tests[which_shard - 1 :: num_shards]

    shards = calculate_shards(num_shards, tests, jobs_to_times)
    _, tests_from_shard = shards[which_shard - 1]
    return tests_from_shard


def get_slow_tests_based_on_S3(test_list: List[str], td_list: List[str], slow_test_threshold: int) -> List[str]:
    """Get list of slow tests based on historic S3 data.
    """
    jobs_to_times: Dict[str, float] = _query_past_job_times()

    # Got no stats from S3, returning early to save runtime
    if len(jobs_to_times) == 0:
        print('Gathered no stats from S3. No new slow tests calculated.')
        return []

    slow_tests: List[str] = []
    for test in test_list:
        if test in jobs_to_times and test not in td_list:
            if jobs_to_times[test] > slow_test_threshold:
                slow_tests.append(test)
    return slow_tests


def get_specified_test_cases(filename: str, tests: List[str]) -> Dict[str, List[str]]:
    """Get test cases from a specified test case file. Usually exported manually or through CI system.
    """
    if not os.path.exists(filename):
        print(f'Could not find specified tests file: {filename}. Proceeding with default behavior.')
        return dict()

    # The below encoding is utf-8-sig because utf-8 doesn't properly handle the byte-order-mark character
    with open(filename, mode='r', encoding="utf-8-sig") as csv_file:
        csv_reader = csv.DictReader(csv_file)
        line_count = 0
        specified_test_case_dict: Dict[str, List[str]] = dict()
        for row in csv_reader:
            line_count += 1
            if line_count == 1:
                if 'test_filename' not in row or 'test_case_name' not in row:
                    print('Data is missing necessary columns for test specification. Proceeding with default behavior.')
                    return dict()
            test_filename = row['test_filename']
            test_case_name = row['test_case_name']
            if test_filename not in tests:
                print(f'Specified test_filename {test_filename} not found in TESTS. Skipping.')
                continue
            if test_filename not in specified_test_case_dict:
                specified_test_case_dict[test_filename] = []
            specified_test_case_dict[test_filename].append(test_case_name)
        print(f'Processed {line_count} test cases.')
        return specified_test_case_dict


def get_reordered_tests(tests: List[str], is_reordering_by_pr: bool) -> List[str]:
    """Get the reordered test filename list based on github PR history or git changed file.
    """
    prioritized_tests = []
    # Try using historic stats from PR.
    if is_reordering_by_pr and HAVE_BOTO3:
        pr_number = os.environ.get("CIRCLE_PR_NUMBER", "")
        if len(pr_number):
            ci_job_prefix = _get_stripped_CI_job()
            s3_reports: List[Tuple["Report", str]] = get_previous_reports_for_pr(
                pr_number, ci_job_prefix)
            prioritized_tests = _query_failure_test_module(s3_reports)
            print("Prioritized test from previous CI info.")

    # Using file changes priority if no stats found from previous PR.
    if len(prioritized_tests) == 0:
        try:
            changed_files = _query_changed_test_files()
        except Exception:
            # If unable to get changed files from git, quit without doing any sorting
            return tests

        prefix = f"test{os.path.sep}"
        prioritized_tests = [f for f in changed_files if f.startswith(prefix) and f.endswith(".py")]
        prioritized_tests = [f[len(prefix):] for f in prioritized_tests]
        prioritized_tests = [f[:-len(".py")] for f in prioritized_tests]
        print("Prioritized test from test file changes.")

    bring_to_front = []
    the_rest = []

    for test in tests:
        if test in prioritized_tests:
            bring_to_front.append(test)
        else:
            the_rest.append(test)
    if len(tests) == len(bring_to_front) + len(the_rest):
        print(f"reordering tests for PR:\n"
              f"prioritized: {bring_to_front}\nthe rest: {the_rest}\n")
        return bring_to_front + the_rest
    else:
        print(f"Something went wrong in CI reordering, expecting total of {len(tests)}:\n"
              f"but found prioritized: {len(bring_to_front)}\nthe rest: {len(the_rest)}\n")
        return tests


# TODO Refactor this and unify with tools.stats.export_slow_tests
def export_S3_test_times(test_times_filename: Optional[str] = None) -> Dict[str, float]:
    test_times: Dict[str, float] = _pull_job_times_from_S3()
    if test_times_filename is not None:
        print(f'Exporting S3 test stats to {test_times_filename}.')
        if os.path.exists(test_times_filename):
            print(f'Overwriting existent file: {test_times_filename}')
        with open(test_times_filename, 'w+') as file:
            job_times_json = _get_job_times_json(test_times)
            json.dump(job_times_json, file, indent='    ', separators=(',', ': '))
            file.write('\n')
    return test_times


def get_test_case_configs(dirpath: str) -> None:
    get_slow_tests(dirpath=dirpath)
    get_disabled_tests(dirpath=dirpath)