1#!/usr/bin/env python3 2 3# Copyright 2020 The gRPC Authors 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16 17# Library to extract scenario definitions from scenario_config.py. 18# 19# Contains functions to filter, analyze and dump scenario definitions. 20# 21# This library is used in loadtest_config.py to generate the "scenariosJSON" 22# field in the format accepted by the OSS benchmarks framework. 23# See https://github.com/grpc/test-infra/blob/master/config/samples/cxx_example_loadtest.yaml 24# 25# It can also be used to dump scenarios to files, to count scenarios by 26# language, and to export scenario languages in a format that can be used for 27# automation. 28# 29# Example usage: 30# 31# scenario_config.py --export_scenarios -l cxx -f cxx_scenario_ -r '.*' \ 32# --category=scalable 33# 34# scenario_config.py --count_scenarios 35# 36# scenario_config.py --count_scenarios --category=scalable 37# 38# For usage of the language config output, see loadtest_config.py. 39 40import argparse 41import collections 42import json 43import re 44import sys 45 46from typing import Any, Callable, Dict, Iterable, NamedTuple 47 48import scenario_config 49 50# Language parameters for load test config generation. 51LanguageConfig = NamedTuple('LanguageConfig', [('category', str), 52 ('language', str), 53 ('client_language', str), 54 ('server_language', str)]) 55 56 57def as_dict_no_empty_values(self): 58 """Returns the parameters as a dictionary, ignoring empty values.""" 59 return dict((item for item in self._asdict().items() if item[1])) 60 61 62def category_string(categories: Iterable[str], category: str) -> str: 63 """Converts a list of categories into a single string for counting.""" 64 if category != 'all': 65 return category if category in categories else '' 66 67 main_categories = ('scalable', 'smoketest') 68 s = set(categories) 69 70 c = [m for m in main_categories if m in s] 71 s.difference_update(main_categories) 72 c.extend(s) 73 return ' '.join(c) 74 75 76def gen_scenario_languages(category: str) -> Iterable[LanguageConfig]: 77 """Generates tuples containing the languages specified in each scenario.""" 78 for language in scenario_config.LANGUAGES: 79 for scenario in scenario_config.LANGUAGES[language].scenarios(): 80 client_language = scenario.get('CLIENT_LANGUAGE', '') 81 server_language = scenario.get('SERVER_LANGUAGE', '') 82 categories = scenario.get('CATEGORIES', []) 83 if category != 'all' and category not in categories: 84 continue 85 cat = category_string(categories, category) 86 yield LanguageConfig(category=cat, 87 language=language, 88 client_language=client_language, 89 server_language=server_language) 90 91 92def scenario_filter( 93 scenario_name_regex: str = '.*', 94 category: str = 'all', 95 client_language: str = '', 96 server_language: str = '', 97) -> Callable[[Dict[str, Any]], bool]: 98 """Returns a function to filter scenarios to process.""" 99 100 def filter_scenario(scenario: Dict[str, Any]) -> bool: 101 """Filters scenarios that match specified criteria.""" 102 if not re.search(scenario_name_regex, scenario["name"]): 103 return False 104 # if the 'CATEGORIES' key is missing, treat scenario as part of 105 # 'scalable' and 'smoketest'. This matches the behavior of 106 # run_performance_tests.py. 107 scenario_categories = scenario.get('CATEGORIES', 108 ['scalable', 'smoketest']) 109 if category not in scenario_categories and category != 'all': 110 return False 111 112 scenario_client_language = scenario.get('CLIENT_LANGUAGE', '') 113 if client_language != scenario_client_language: 114 return False 115 116 scenario_server_language = scenario.get('SERVER_LANGUAGE', '') 117 if server_language != scenario_server_language: 118 return False 119 120 return True 121 122 return filter_scenario 123 124 125def gen_scenarios( 126 language_name: str, scenario_filter_function: Callable[[Dict[str, Any]], 127 bool] 128) -> Iterable[Dict[str, Any]]: 129 """Generates scenarios that match a given filter function.""" 130 return map( 131 scenario_config.remove_nonproto_fields, 132 filter(scenario_filter_function, 133 scenario_config.LANGUAGES[language_name].scenarios())) 134 135 136def dump_to_json_files(scenarios: Iterable[Dict[str, Any]], 137 filename_prefix: str) -> None: 138 """Dumps a list of scenarios to JSON files""" 139 count = 0 140 for scenario in scenarios: 141 filename = '{}{}.json'.format(filename_prefix, scenario['name']) 142 print('Writing file {}'.format(filename), file=sys.stderr) 143 with open(filename, 'w') as outfile: 144 # The dump file should have {"scenarios" : []} as the top level 145 # element, when embedded in a LoadTest configuration YAML file. 146 json.dump({'scenarios': [scenario]}, outfile, indent=2) 147 count += 1 148 print('Wrote {} scenarios'.format(count), file=sys.stderr) 149 150 151def main() -> None: 152 language_choices = sorted(scenario_config.LANGUAGES.keys()) 153 argp = argparse.ArgumentParser(description='Exports scenarios to files.') 154 argp.add_argument('--export_scenarios', 155 action='store_true', 156 help='Export scenarios to JSON files.') 157 argp.add_argument('--count_scenarios', 158 action='store_true', 159 help='Count scenarios for all test languages.') 160 argp.add_argument('-l', 161 '--language', 162 choices=language_choices, 163 help='Language to export.') 164 argp.add_argument('-f', 165 '--filename_prefix', 166 default='scenario_dump_', 167 type=str, 168 help='Prefix for exported JSON file names.') 169 argp.add_argument('-r', 170 '--regex', 171 default='.*', 172 type=str, 173 help='Regex to select scenarios to run.') 174 argp.add_argument( 175 '--category', 176 default='all', 177 choices=['all', 'inproc', 'scalable', 'smoketest', 'sweep'], 178 help='Select scenarios for a category of tests.') 179 argp.add_argument( 180 '--client_language', 181 default='', 182 choices=language_choices, 183 help='Select only scenarios with a specified client language.') 184 argp.add_argument( 185 '--server_language', 186 default='', 187 choices=language_choices, 188 help='Select only scenarios with a specified server language.') 189 args = argp.parse_args() 190 191 if args.export_scenarios and not args.language: 192 print('Dumping scenarios requires a specified language.', 193 file=sys.stderr) 194 argp.print_usage(file=sys.stderr) 195 return 196 197 if args.export_scenarios: 198 s_filter = scenario_filter(scenario_name_regex=args.regex, 199 category=args.category, 200 client_language=args.client_language, 201 server_language=args.server_language) 202 scenarios = gen_scenarios(args.language, s_filter) 203 dump_to_json_files(scenarios, args.filename_prefix) 204 205 if args.count_scenarios: 206 print('Scenario count for all languages (category: {}):'.format( 207 args.category)) 208 print('{:>5} {:16} {:8} {:8} {}'.format('Count', 'Language', 'Client', 209 'Server', 'Categories')) 210 c = collections.Counter(gen_scenario_languages(args.category)) 211 total = 0 212 for ((cat, l, cl, sl), count) in c.most_common(): 213 print('{count:5} {l:16} {cl:8} {sl:8} {cat}'.format(l=l, 214 cl=cl, 215 sl=sl, 216 count=count, 217 cat=cat)) 218 total += count 219 220 print('\n{:>5} total scenarios (category: {})'.format( 221 total, args.category)) 222 223 224if __name__ == "__main__": 225 main() 226