modelscope--ms-swift/tests/run.py

#!/usr/bin/env python
# Copyright (c) ModelScope Contributors. All rights reserved.

import argparse
import datetime
import math
import os
import pandas
import subprocess
import sys
import tempfile
import time
import torch
import unittest
import yaml
from fnmatch import fnmatch
from model_tag import ModelTag, commit_model_ut_result
from pathlib import Path
from test_utils import get_case_model_info
from unittest import TextTestResult

from swift.utils import get_logger

logger = get_logger()


def deduplicate_preserve_order(items):
    return list(dict.fromkeys(items))


def get_available_npu_devices(visible_npus):
    npu_devices = [device.strip() for device in visible_npus.split(',') if device.strip()]
    if not npu_devices:
        return []

    try:
        import torch_npu  # noqa: F401
        npu_count = torch.npu.device_count() if hasattr(torch, 'npu') and torch.npu.is_available() else 0
    except Exception as e:
        logger.warning('Failed to query torch.npu.device_count(): %s' % e)
        return []

    if npu_count <= 0:
        logger.warning('ASCEND_RT_VISIBLE_DEVICES=%s, but torch.npu.device_count()=%s' % (visible_npus, npu_count))
        return []
    if npu_count < len(npu_devices):
        logger.warning('ASCEND_RT_VISIBLE_DEVICES=%s, but torch.npu.device_count()=%s; using %s' %
                       (visible_npus, npu_count, ','.join(npu_devices[:npu_count])))
    return npu_devices[:npu_count]


def test_cases_result_to_df(result_list):
    table_header = ['Name', 'Result', 'Info', 'Start time', 'Stop time', 'Time cost(seconds)']
    df = pandas.DataFrame(result_list, columns=table_header).sort_values(by=['Start time'], ascending=True)
    return df


def statistics_test_result(df):
    total_cases = df.shape[0]
    # yapf: disable
    success_cases = df.loc[df['Result'] == 'Success'].shape[0]
    error_cases = df.loc[df['Result'] == 'Error'].shape[0]
    failures_cases = df.loc[df['Result'] == 'Failures'].shape[0]
    expected_failure_cases = df.loc[df['Result'] == 'ExpectedFailures'].shape[0]
    unexpected_success_cases = df.loc[df['Result'] == 'UnexpectedSuccesses'].shape[0]
    skipped_cases = df.loc[df['Result'] == 'Skipped'].shape[0]
    # yapf: enable

    if failures_cases > 0 or \
       error_cases > 0 or \
       unexpected_success_cases > 0:
        final_result = 'FAILED'
    else:
        final_result = 'SUCCESS'
    result_msg = '%s (Runs=%s,success=%s,failures=%s,errors=%s,\
    skipped=%s,expected failures=%s,unexpected successes=%s)' % (final_result, total_cases, success_cases,
                                                                 failures_cases, error_cases, skipped_cases,
                                                                 expected_failure_cases, unexpected_success_cases)

    model_cases = get_case_model_info()
    for model_name, case_info in model_cases.items():
        cases = df.loc[df['Name'].str.contains('|'.join(list(case_info)))]
        results = cases['Result']
        result = None
        if any(results == 'Error') or any(results == 'Failures') or any(results == 'UnexpectedSuccesses'):
            result = ModelTag.MODEL_FAIL
        elif any(results == 'Success'):
            result = ModelTag.MODEL_PASS
        elif all(results == 'Skipped'):
            result = ModelTag.MODEL_SKIP
        else:
            print(f'invalid results for {model_name} \n{result}')

        if result is not None:
            commit_model_ut_result(model_name, result)
    print('Testing result summary.')
    print(result_msg)
    if final_result == 'FAILED':
        sys.exit(1)


def gather_test_suites_in_files(test_dir, case_file_list, list_tests):
    test_suite = unittest.TestSuite()
    for case in deduplicate_preserve_order(case_file_list):
        test_case = unittest.defaultTestLoader.discover(start_dir=test_dir, pattern=case)
        test_suite.addTest(test_case)
        if hasattr(test_case, '__iter__'):
            for subcase in test_case:
                if list_tests:
                    print(subcase)
        else:
            if list_tests:
                print(test_case)
    return test_suite


def gather_test_suites_files(test_dir, pattern):
    case_file_list = []
    for dirpath, dirnames, filenames in os.walk(test_dir):
        for file in filenames:
            if fnmatch(file, pattern):
                case_file_list.append(file)

    return deduplicate_preserve_order(case_file_list)


def collect_test_results(case_results):
    result_list = []  # each item is Case, Result, Start time, Stop time, Time cost
    for case_result in case_results.successes:
        result_list.append((case_result.test_full_name, 'Success', '', case_result.start_time, case_result.stop_time,
                            case_result.time_cost))
    for case_result in case_results.errors:
        result_list.append((case_result[0].test_full_name, 'Error', case_result[1], case_result[0].start_time,
                            case_result[0].stop_time, case_result[0].time_cost))
    for case_result in case_results.skipped:
        result_list.append((case_result[0].test_full_name, 'Skipped', case_result[1], case_result[0].start_time,
                            case_result[0].stop_time, case_result[0].time_cost))
    for case_result in case_results.expectedFailures:
        result_list.append((case_result[0].test_full_name, 'ExpectedFailures', case_result[1],
                            case_result[0].start_time, case_result[0].stop_time, case_result[0].time_cost))
    for case_result in case_results.failures:
        result_list.append((case_result[0].test_full_name, 'Failures', case_result[1], case_result[0].start_time,
                            case_result[0].stop_time, case_result[0].time_cost))
    for case_result in case_results.unexpectedSuccesses:
        result_list.append((case_result.test_full_name, 'UnexpectedSuccesses', '', case_result.start_time,
                            case_result.stop_time, case_result.time_cost))
    return result_list


def run_command_with_popen(cmd):
    with subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, encoding='utf8') as sub_process:
        for line in iter(sub_process.stdout.readline, ''):
            sys.stdout.write(line)


def async_run_command_with_popen(cmd, device_id):
    logger.info('Worker id: %s args: %s' % (device_id, cmd))
    env = os.environ.copy()
    visible_npus = env.get('ASCEND_RT_VISIBLE_DEVICES')
    if visible_npus:
        npu_devices = get_available_npu_devices(visible_npus)
        if npu_devices:
            env['ASCEND_RT_VISIBLE_DEVICES'] = npu_devices[device_id % len(npu_devices)]
            logger.info('Worker id: %s ASCEND_RT_VISIBLE_DEVICES: %s' % (device_id, env['ASCEND_RT_VISIBLE_DEVICES']))
    else:
        env['CUDA_VISIBLE_DEVICES'] = '%s' % device_id
    sub_process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
        universal_newlines=True,
        env=env,
        encoding='utf8')
    return sub_process


def save_test_result(df, args):
    if args.result_dir is not None:
        file_name = str(int(datetime.datetime.now().timestamp() * 1000))
        os.umask(0)
        Path(args.result_dir).mkdir(mode=0o777, parents=True, exist_ok=True)
        Path(os.path.join(args.result_dir, file_name)).touch(mode=0o666, exist_ok=True)
        df.to_pickle(os.path.join(args.result_dir, file_name))


def run_command(cmd):
    logger.info('Running command: %s' % ' '.join(cmd))
    response = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        response.check_returncode()
        logger.info(response.stdout.decode('utf8'))
    except subprocess.CalledProcessError as error:
        logger.error('stdout: %s, stderr: %s' % (response.stdout.decode('utf8'), error.stderr.decode('utf8')))


def install_packages(pkgs):
    if pkgs is None:
        return
    cmd = [sys.executable, '-m', 'pip', 'install']
    for pkg in pkgs:
        cmd.append(pkg)

    run_command(cmd)


def install_requirements(requirements):
    for req in requirements:
        cmd = [
            sys.executable, '-m', 'pip', 'install', '-r',
            'requirements/%s' % req, '-f', 'https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html'
        ]
        run_command(cmd)


def wait_for_free_worker(workers):
    while True:
        for idx, worker in enumerate(workers):
            if worker is None:
                logger.info('return free worker: %s' % (idx))
                return idx
            if worker.poll() is None:  # running, get output
                for line in iter(worker.stdout.readline, ''):
                    if line != '':
                        sys.stdout.write(line)
                    else:
                        break
            else:  # worker process completed.
                logger.info('Process end: %s' % (idx))
                workers[idx] = None
                return idx
        time.sleep(0.001)


def wait_for_workers(workers):
    while True:
        for idx, worker in enumerate(workers):
            if worker is None:
                continue
            # check worker is completed.
            if worker.poll() is None:
                for line in iter(worker.stdout.readline, ''):
                    if line != '':
                        sys.stdout.write(line)
                    else:
                        break
            else:
                logger.info('Process idx: %s end!' % (idx))
                workers[idx] = None

        is_all_completed = True
        for idx, worker in enumerate(workers):
            if worker is not None:
                is_all_completed = False
                break

        if is_all_completed:
            logger.info('All sub process is completed!')
            break
        time.sleep(0.001)


def parallel_run_case_in_env(env_name, env, test_suite_env_map, isolated_cases, result_dir, parallel):
    logger.info('Running case in env: %s' % env_name)
    # install requirements and deps # run_config['envs'][env]
    if 'requirements' in env:
        install_requirements(env['requirements'])
    if 'dependencies' in env:
        install_packages(env['dependencies'])
    # case worker processes
    worker_processes = [None] * parallel
    for test_suite_file in isolated_cases:  # run case in subprocess
        if test_suite_file in test_suite_env_map and test_suite_env_map[test_suite_file] == env_name:
            cmd = [
                'python',
                'tests/run.py',
                '--pattern',
                test_suite_file,
                '--result_dir',
                result_dir,
            ]
            worker_idx = wait_for_free_worker(worker_processes)
            worker_process = async_run_command_with_popen(cmd, worker_idx)
            os.set_blocking(worker_process.stdout.fileno(), False)
            worker_processes[worker_idx] = worker_process
        else:
            pass  # case not in run list.

    # run remain cases in a process.
    remain_suite_files = []
    for k, v in test_suite_env_map.items():
        if k not in isolated_cases and v == env_name:
            remain_suite_files.append(k)
    if len(remain_suite_files) == 0:
        wait_for_workers(worker_processes)
        return
    # roughly split case in parallel
    part_count = math.ceil(len(remain_suite_files) / parallel)
    suites_chunks = [remain_suite_files[x:x + part_count] for x in range(0, len(remain_suite_files), part_count)]
    for suites_chunk in suites_chunks:
        worker_idx = wait_for_free_worker(worker_processes)
        cmd = ['python', 'tests/run.py', '--result_dir', result_dir, '--suites']
        for suite in suites_chunk:
            cmd.append(suite)
        worker_process = async_run_command_with_popen(cmd, worker_idx)
        os.set_blocking(worker_process.stdout.fileno(), False)
        worker_processes[worker_idx] = worker_process

    wait_for_workers(worker_processes)


def run_case_in_env(env_name, env, test_suite_env_map, isolated_cases, result_dir):
    # install requirements and deps # run_config['envs'][env]
    if 'requirements' in env:
        install_requirements(env['requirements'])
    if 'dependencies' in env:
        install_packages(env['dependencies'])

    for test_suite_file in isolated_cases:  # run case in subprocess
        if test_suite_file in test_suite_env_map and test_suite_env_map[test_suite_file] == env_name:
            cmd = [
                'python',
                'tests/run.py',
                '--pattern',
                test_suite_file,
                '--result_dir',
                result_dir,
            ]
            run_command_with_popen(cmd)
        else:
            pass  # case not in run list.

    # run remain cases in a process.
    remain_suite_files = []
    for k, v in test_suite_env_map.items():
        if k not in isolated_cases and v == env_name:
            remain_suite_files.append(k)
    if len(remain_suite_files) == 0:
        return
    cmd = ['python', 'tests/run.py', '--result_dir', result_dir, '--suites']
    for suite in remain_suite_files:
        cmd.append(suite)
    run_command_with_popen(cmd)


def run_non_parallelizable_test_suites(suites, result_dir):
    if len(suites) == 0:
        return
    cmd = ['python', 'tests/run.py', '--result_dir', result_dir, '--suites']
    for suite in suites:
        cmd.append(suite)
    run_command_with_popen(cmd)


# Selected cases:
def get_selected_cases():
    cmd = ['python', '-u', 'tests/run_analysis.py']
    selected_cases = []
    with subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, encoding='utf8') as sub_process:
        for line in iter(sub_process.stdout.readline, ''):
            sys.stdout.write(line)
            if line.startswith('Selected cases:'):
                line = line.replace('Selected cases:', '').strip()
                selected_cases = line.split(',')
        sub_process.wait()
        if sub_process.returncode != 0:
            msg = 'Run analysis exception, returncode: %s!' % sub_process.returncode
            logger.error(msg)
            raise Exception(msg)
    return selected_cases


def run_in_subprocess(args):
    # only case args.isolated_cases run in subprocess, all other run in a subprocess
    if not args.no_diff:  # run based on git diff
        try:
            test_suite_files = get_selected_cases()
            logger.info('Tests suite to run: ')
            for f in test_suite_files:
                logger.info(f)
        except Exception:
            logger.error('Get test suite based diff exception!, will run all cases.')
            test_suite_files = gather_test_suites_files(os.path.abspath(args.test_dir), args.pattern)
        if len(test_suite_files) == 0:
            logger.error('Get no test suite based on diff, run all the cases.')
            test_suite_files = gather_test_suites_files(os.path.abspath(args.test_dir), args.pattern)
    else:
        test_suite_files = gather_test_suites_files(os.path.abspath(args.test_dir), args.pattern)
    test_suite_files = deduplicate_preserve_order(test_suite_files)

    non_parallelizable_suites = []
    test_suite_files = [x for x in test_suite_files if x not in non_parallelizable_suites]

    run_config = None
    isolated_cases = []
    test_suite_env_map = {}
    # put all the case in default env.
    for test_suite_file in test_suite_files:
        test_suite_env_map[test_suite_file] = 'default'

    if args.run_config is not None and Path(args.run_config).exists():
        with open(args.run_config, encoding='utf-8') as f:
            run_config = yaml.safe_load(f)
        if 'isolated' in run_config:
            isolated_cases = run_config['isolated']

        if 'envs' in run_config:
            for env in run_config['envs']:
                if env != 'default':
                    for test_suite in run_config['envs'][env]['tests']:
                        if test_suite in test_suite_env_map:
                            test_suite_env_map[test_suite] = env

    if args.subprocess:  # run all case in subprocess
        isolated_cases = test_suite_files

    with tempfile.TemporaryDirectory() as temp_result_dir:
        # first run cases that nonparallelizable
        run_non_parallelizable_test_suites(non_parallelizable_suites, temp_result_dir)

        # run case parallel in envs
        for env in set(test_suite_env_map.values()):
            parallel_run_case_in_env(env, run_config['envs'][env], test_suite_env_map, isolated_cases, temp_result_dir,
                                     args.parallel)

        result_dfs = []
        result_path = Path(temp_result_dir)
        for result in result_path.iterdir():
            if Path.is_file(result):
                df = pandas.read_pickle(result)
                result_dfs.append(df)
        result_pd = pandas.concat(result_dfs)  # merge result of every test suite.
        print_table_result(result_pd)
        print_abnormal_case_info(result_pd)
        statistics_test_result(result_pd)


def get_object_full_name(obj):
    klass = obj.__class__
    module = klass.__module__
    if module == 'builtins':
        return klass.__qualname__
    return module + '.' + klass.__qualname__


class TimeCostTextTestResult(TextTestResult):
    """Record test case time used!"""

    def __init__(self, stream, descriptions, verbosity):
        self.successes = []
        super(TimeCostTextTestResult, self).__init__(stream, descriptions, verbosity)

    def startTest(self, test):
        test.start_time = datetime.datetime.now()
        test.test_full_name = get_object_full_name(test) + '.' + test._testMethodName
        self.stream.writeln('Test case:  %s start at: %s' % (test.test_full_name, test.start_time))

        return super(TimeCostTextTestResult, self).startTest(test)

    def stopTest(self, test):
        TextTestResult.stopTest(self, test)
        test.stop_time = datetime.datetime.now()
        test.time_cost = (test.stop_time - test.start_time).total_seconds()
        self.stream.writeln('Test case: %s stop at: %s, cost time: %s(seconds)' %
                            (test.test_full_name, test.stop_time, test.time_cost))
        if torch.cuda.is_available() and test.time_cost > 5.0:  # print nvidia-smi
            cmd = ['nvidia-smi']
            run_command_with_popen(cmd)
        super(TimeCostTextTestResult, self).stopTest(test)

    def addSuccess(self, test):
        self.successes.append(test)
        super(TextTestResult, self).addSuccess(test)


class TimeCostTextTestRunner(unittest.runner.TextTestRunner):
    resultclass = TimeCostTextTestResult

    def run(self, test):
        return super(TimeCostTextTestRunner, self).run(test)

    def _makeResult(self):
        result = super(TimeCostTextTestRunner, self)._makeResult()
        return result


def gather_test_cases(test_dir, pattern, list_tests):
    case_list = gather_test_suites_files(test_dir, pattern)

    test_suite = unittest.TestSuite()

    for case in case_list:
        test_case = unittest.defaultTestLoader.discover(start_dir=test_dir, pattern=case)
        test_suite.addTest(test_case)
        if hasattr(test_case, '__iter__'):
            for subcase in test_case:
                if list_tests:
                    print(subcase)
        else:
            if list_tests:
                print(test_case)
    return test_suite


def print_abnormal_case_info(df):
    df = df.loc[(df['Result'] == 'Error') | (df['Result'] == 'Failures')]
    for _, row in df.iterrows():
        print('Case %s run result: %s, msg:\n%s' % (row['Name'], row['Result'], row['Info']))


def print_table_result(df):
    df = df.loc[df['Result'] != 'Skipped']
    df = df.drop('Info', axis=1)
    formatters = {
        'Name': '{{:<{}s}}'.format(df['Name'].str.len().max()).format,
        'Result': '{{:<{}s}}'.format(df['Result'].str.len().max()).format,
    }
    with pandas.option_context('display.max_rows', None, 'display.max_columns', None, 'display.width', None):
        print(df.to_string(justify='left', formatters=formatters, index=False))


def main(args):
    runner = TimeCostTextTestRunner()
    if args.suites is not None and len(args.suites) > 0:
        logger.info('Running: %s' % ' '.join(args.suites))
        test_suite = gather_test_suites_in_files(args.test_dir, args.suites, args.list_tests)
    else:
        test_suite = gather_test_cases(os.path.abspath(args.test_dir), args.pattern, args.list_tests)
    if not args.list_tests:
        result = runner.run(test_suite)
        logger.info('Running case completed, pid: %s, suites: %s' % (os.getpid(), args.suites))
        result = collect_test_results(result)
        df = test_cases_result_to_df(result)
        if args.result_dir is not None:
            save_test_result(df, args)
        else:
            print_table_result(df)
            print_abnormal_case_info(df)
            statistics_test_result(df)


if __name__ == '__main__':
    parser = argparse.ArgumentParser('test runner')
    parser.add_argument('--list_tests', action='store_true', help='list all tests')
    parser.add_argument('--pattern', default='test_*.py', help='test file pattern')
    parser.add_argument('--test_dir', default='tests', help='directory to be tested')
    parser.add_argument('--level', default=0, type=int, help='2 -- all, 1 -- p1, 0 -- p0')
    parser.add_argument('--profile', action='store_true', help='enable profiling')
    parser.add_argument('--run_config', default=None, help='specified case run config file(yaml file)')
    parser.add_argument('--subprocess', action='store_true', help='run all test suite in subprocess')
    parser.add_argument('--result_dir', default=None, help='Save result to directory, internal use only')
    parser.add_argument(
        '--parallel', default=1, type=int, help='Set case parallels, default single process, set with gpu number.')
    parser.add_argument(
        '--no-diff',
        action='store_true',
        help='Default running case based on git diff(with master), disable with --no-diff)')
    parser.add_argument('--suites', nargs='*', help='Run specified test suites(test suite files list split by space)')
    args = parser.parse_args()
    print(args)
    if args.run_config is not None or args.subprocess:
        run_in_subprocess(args)
    else:
        main(args)