grpc1.45.2/tools/run_tests/performance/scenario_config_exporter.py

#!/usr/bin/env python3

# Copyright 2020 The gRPC Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Library to extract scenario definitions from scenario_config.py.
#
# Contains functions to filter, analyze and dump scenario definitions.
#
# This library is used in loadtest_config.py to generate the "scenariosJSON"
# field in the format accepted by the OSS benchmarks framework.
# See https://github.com/grpc/test-infra/blob/master/config/samples/cxx_example_loadtest.yaml
#
# It can also be used to dump scenarios to files, to count scenarios by
# language, and to export scenario languages in a format that can be used for
# automation.
#
# Example usage:
#
# scenario_config.py --export_scenarios -l cxx -f cxx_scenario_ -r '.*' \
#     --category=scalable
#
# scenario_config.py --count_scenarios
#
# scenario_config.py --count_scenarios --category=scalable
#
# For usage of the language config output, see loadtest_config.py.

import argparse
import collections
import json
import re
import sys
from typing import Any, Callable, Dict, Iterable, NamedTuple

import scenario_config

# Language parameters for load test config generation.

LanguageConfig = NamedTuple('LanguageConfig', [('category', str),
                                               ('language', str),
                                               ('client_language', str),
                                               ('server_language', str)])


def category_string(categories: Iterable[str], category: str) -> str:
    """Converts a list of categories into a single string for counting."""
    if category != 'all':
        return category if category in categories else ''

    main_categories = ('scalable', 'smoketest')
    s = set(categories)

    c = [m for m in main_categories if m in s]
    s.difference_update(main_categories)
    c.extend(s)
    return ' '.join(c)


def gen_scenario_languages(category: str) -> Iterable[LanguageConfig]:
    """Generates tuples containing the languages specified in each scenario."""
    for language in scenario_config.LANGUAGES:
        for scenario in scenario_config.LANGUAGES[language].scenarios():
            client_language = scenario.get('CLIENT_LANGUAGE', '')
            server_language = scenario.get('SERVER_LANGUAGE', '')
            categories = scenario.get('CATEGORIES', [])
            if category != 'all' and category not in categories:
                continue
            cat = category_string(categories, category)
            yield LanguageConfig(category=cat,
                                 language=language,
                                 client_language=client_language,
                                 server_language=server_language)


def scenario_filter(
    scenario_name_regex: str = '.*',
    category: str = 'all',
    client_language: str = '',
    server_language: str = '',
) -> Callable[[Dict[str, Any]], bool]:
    """Returns a function to filter scenarios to process."""

    def filter_scenario(scenario: Dict[str, Any]) -> bool:
        """Filters scenarios that match specified criteria."""
        if not re.search(scenario_name_regex, scenario["name"]):
            return False
        # if the 'CATEGORIES' key is missing, treat scenario as part of
        # 'scalable' and 'smoketest'. This matches the behavior of
        # run_performance_tests.py.
        scenario_categories = scenario.get('CATEGORIES',
                                           ['scalable', 'smoketest'])
        if category not in scenario_categories and category != 'all':
            return False

        scenario_client_language = scenario.get('CLIENT_LANGUAGE', '')
        if client_language != scenario_client_language:
            return False

        scenario_server_language = scenario.get('SERVER_LANGUAGE', '')
        if server_language != scenario_server_language:
            return False

        return True

    return filter_scenario


def gen_scenarios(
    language_name: str, scenario_filter_function: Callable[[Dict[str, Any]],
                                                           bool]
) -> Iterable[Dict[str, Any]]:
    """Generates scenarios that match a given filter function."""
    return map(
        scenario_config.remove_nonproto_fields,
        filter(scenario_filter_function,
               scenario_config.LANGUAGES[language_name].scenarios()))


def dump_to_json_files(scenarios: Iterable[Dict[str, Any]],
                       filename_prefix: str) -> None:
    """Dumps a list of scenarios to JSON files"""
    count = 0
    for scenario in scenarios:
        filename = '{}{}.json'.format(filename_prefix, scenario['name'])
        print('Writing file {}'.format(filename), file=sys.stderr)
        with open(filename, 'w') as outfile:
            # The dump file should have {"scenarios" : []} as the top level
            # element, when embedded in a LoadTest configuration YAML file.
            json.dump({'scenarios': [scenario]}, outfile, indent=2)
        count += 1
    print('Wrote {} scenarios'.format(count), file=sys.stderr)


def main() -> None:
    language_choices = sorted(scenario_config.LANGUAGES.keys())
    argp = argparse.ArgumentParser(description='Exports scenarios to files.')
    argp.add_argument('--export_scenarios',
                      action='store_true',
                      help='Export scenarios to JSON files.')
    argp.add_argument('--count_scenarios',
                      action='store_true',
                      help='Count scenarios for all test languages.')
    argp.add_argument('-l',
                      '--language',
                      choices=language_choices,
                      help='Language to export.')
    argp.add_argument('-f',
                      '--filename_prefix',
                      default='scenario_dump_',
                      type=str,
                      help='Prefix for exported JSON file names.')
    argp.add_argument('-r',
                      '--regex',
                      default='.*',
                      type=str,
                      help='Regex to select scenarios to run.')
    argp.add_argument(
        '--category',
        default='all',
        choices=['all', 'inproc', 'scalable', 'smoketest', 'sweep'],
        help='Select scenarios for a category of tests.')
    argp.add_argument(
        '--client_language',
        default='',
        choices=language_choices,
        help='Select only scenarios with a specified client language.')
    argp.add_argument(
        '--server_language',
        default='',
        choices=language_choices,
        help='Select only scenarios with a specified server language.')
    args = argp.parse_args()

    if args.export_scenarios and not args.language:
        print('Dumping scenarios requires a specified language.',
              file=sys.stderr)
        argp.print_usage(file=sys.stderr)
        return

    if args.export_scenarios:
        s_filter = scenario_filter(scenario_name_regex=args.regex,
                                   category=args.category,
                                   client_language=args.client_language,
                                   server_language=args.server_language)
        scenarios = gen_scenarios(args.language, s_filter)
        dump_to_json_files(scenarios, args.filename_prefix)

    if args.count_scenarios:
        print('Scenario count for all languages (category: {}):'.format(
            args.category))
        print('{:>5}  {:16} {:8} {:8} {}'.format('Count', 'Language', 'Client',
                                                 'Server', 'Categories'))
        c = collections.Counter(gen_scenario_languages(args.category))
        total = 0
        for ((cat, l, cl, sl), count) in c.most_common():
            print('{count:5}  {l:16} {cl:8} {sl:8} {cat}'.format(l=l,
                                                                 cl=cl,
                                                                 sl=sl,
                                                                 count=count,
                                                                 cat=cat))
            total += count

        print('\n{:>5}  total scenarios (category: {})'.format(
            total, args.category))


if __name__ == "__main__":
    main()