| #!/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() |
