1#!/usr/bin/env python3 2# Copyright 2021 The gRPC Authors 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16# Script to generate test configurations for the OSS benchmarks framework. 17# 18# This script filters test scenarios and generates uniquely named configurations 19# for each test. Configurations are dumped in multipart YAML format. 20# 21# See documentation below: 22# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md#grpc-oss-benchmarks 23 24import argparse 25import copy 26import datetime 27import itertools 28import os 29import string 30import sys 31import uuid 32 33from typing import Any, Dict, Iterable, Mapping, Optional, Type 34 35import json 36import yaml 37 38import scenario_config 39import scenario_config_exporter 40 41CONFIGURATION_FILE_HEADER_COMMENT = """ 42# Load test configurations generated from a template by loadtest_config.py. 43# See documentation below: 44# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md#grpc-oss-benchmarks 45""" 46 47# TODO(paulosjca): Merge label_language and image_language into one function. 48# These functions are necessary because 'c++' is not allowed as a label value in 49# kubernetes, and because languages share images in the existing templates. Once 50# the templates are reorganized and most image mapping is removed, the two 51# functions can be merged into one. 52 53 54def label_language(language: str) -> str: 55 """Convert scenario language to place in a resource label.""" 56 return { 57 'c++': 'cxx', 58 }.get(language, language) 59 60 61def image_language(language: str) -> str: 62 """Convert scenario languages to image languages.""" 63 return { 64 'c++': 'cxx', 65 'node_purejs': 'node', 66 'php7': 'php', 67 'php7_protobuf_c': 'php', 68 'python_asyncio': 'python', 69 }.get(language, language) 70 71 72def default_prefix() -> str: 73 """Constructs and returns a default prefix for LoadTest names.""" 74 return os.environ.get('USER', 'loadtest') 75 76 77def now_string() -> str: 78 """Returns the current date and time in string format.""" 79 return datetime.datetime.now().strftime('%Y%m%d%H%M%S') 80 81 82def validate_loadtest_name(name: str) -> None: 83 """Validates that a LoadTest name is in the expected format.""" 84 if len(name) > 63: 85 raise ValueError( 86 'LoadTest name must be less than 63 characters long: %s' % name) 87 if not all((s.isalnum() for s in name.split('-'))): 88 raise ValueError('Invalid elements in LoadTest name: %s' % name) 89 90 91def loadtest_base_name(scenario_name: str, 92 uniquifier_elements: Iterable[str]) -> str: 93 """Constructs and returns the base name for a LoadTest resource.""" 94 name_elements = scenario_name.split('_') 95 name_elements.extend(uniquifier_elements) 96 return '-'.join(name_elements) 97 98 99def loadtest_name(prefix: str, scenario_name: str, 100 uniquifier_elements: Iterable[str]) -> str: 101 """Constructs and returns a valid name for a LoadTest resource.""" 102 base_name = loadtest_base_name(scenario_name, uniquifier_elements) 103 name_elements = [] 104 if prefix: 105 name_elements.append(prefix) 106 name_elements.append(str(uuid.uuid5(uuid.NAMESPACE_DNS, base_name))) 107 name = '-'.join(name_elements) 108 validate_loadtest_name(name) 109 return name 110 111 112def validate_annotations(annotations: Dict[str, str]) -> None: 113 """Validates that annotations do not contain reserved names. 114 115 These names are automatically added by the config generator. 116 """ 117 names = set(('scenario', 'uniquifier')).intersection(annotations) 118 if names: 119 raise ValueError('Annotations contain reserved names: %s' % names) 120 121 122def gen_run_indices(runs_per_test: int) -> Iterable[str]: 123 """Generates run indices for multiple runs, as formatted strings.""" 124 if runs_per_test < 2: 125 yield '' 126 return 127 prefix_length = len('{:d}'.format(runs_per_test - 1)) 128 prefix_fmt = '{{:{:d}d}}'.format(prefix_length) 129 for i in range(runs_per_test): 130 yield prefix_fmt.format(i) 131 132 133def gen_loadtest_configs( 134 base_config: Mapping[str, Any], 135 base_config_clients: Iterable[Mapping[str, Any]], 136 base_config_servers: Iterable[Mapping[str, Any]], 137 scenario_name_regex: str, 138 language_config: scenario_config_exporter.LanguageConfig, 139 loadtest_name_prefix: str, 140 uniquifier_elements: Iterable[str], 141 annotations: Mapping[str, str], 142 runs_per_test: int = 1) -> Iterable[Dict[str, Any]]: 143 """Generates LoadTest configurations for a given language config. 144 145 The LoadTest configurations are generated as YAML objects. 146 """ 147 validate_annotations(annotations) 148 prefix = loadtest_name_prefix or default_prefix() 149 cl = image_language(language_config.client_language or 150 language_config.language) 151 sl = image_language(language_config.server_language or 152 language_config.language) 153 scenario_filter = scenario_config_exporter.scenario_filter( 154 scenario_name_regex=scenario_name_regex, 155 category=language_config.category, 156 client_language=language_config.client_language, 157 server_language=language_config.server_language) 158 scenarios = scenario_config_exporter.gen_scenarios(language_config.language, 159 scenario_filter) 160 161 for scenario in scenarios: 162 for run_index in gen_run_indices(runs_per_test): 163 uniq = (uniquifier_elements + 164 [run_index] if run_index else uniquifier_elements) 165 name = loadtest_name(prefix, scenario['name'], uniq) 166 scenario_str = json.dumps({'scenarios': scenario}, 167 indent=' ') + '\n' 168 169 config = copy.deepcopy(base_config) 170 171 metadata = config['metadata'] 172 metadata['name'] = name 173 if 'labels' not in metadata: 174 metadata['labels'] = dict() 175 metadata['labels']['language'] = label_language( 176 language_config.language) 177 metadata['labels']['prefix'] = prefix 178 if 'annotations' not in metadata: 179 metadata['annotations'] = dict() 180 metadata['annotations'].update(annotations) 181 metadata['annotations'].update({ 182 'scenario': scenario['name'], 183 'uniquifier': '-'.join(uniq), 184 }) 185 186 spec = config['spec'] 187 188 # Select clients with the required language. 189 spec['clients'] = [ 190 client for client in base_config_clients 191 if client['language'] == cl 192 ] 193 if not spec['clients']: 194 raise IndexError('Client language not found in template: %s' % 195 cl) 196 197 # Select servers with the required language. 198 spec['servers'] = [ 199 server for server in base_config_servers 200 if server['language'] == sl 201 ] 202 if not spec['servers']: 203 raise IndexError('Server language not found in template: %s' % 204 sl) 205 206 spec['scenariosJSON'] = scenario_str 207 208 yield config 209 210 211def parse_key_value_args(args: Optional[Iterable[str]]) -> Dict[str, str]: 212 """Parses arguments in the form key=value into a dictionary.""" 213 d = dict() 214 if args is None: 215 return d 216 for arg in args: 217 key, equals, value = arg.partition('=') 218 if equals != '=': 219 raise ValueError('Expected key=value: ' + value) 220 d[key] = value 221 return d 222 223 224def config_dumper(header_comment: str) -> Type[yaml.SafeDumper]: 225 """Returns a custom dumper to dump configurations in the expected format.""" 226 227 class ConfigDumper(yaml.SafeDumper): 228 229 def expect_stream_start(self): 230 super().expect_stream_start() 231 if isinstance(self.event, yaml.StreamStartEvent): 232 self.write_indent() 233 self.write_indicator(header_comment, need_whitespace=False) 234 235 def expect_block_sequence(self): 236 super().expect_block_sequence() 237 self.increase_indent() 238 239 def expect_block_sequence_item(self, first=False): 240 if isinstance(self.event, yaml.SequenceEndEvent): 241 self.indent = self.indents.pop() 242 super().expect_block_sequence_item(first) 243 244 def str_presenter(dumper, data): 245 if '\n' in data: 246 return dumper.represent_scalar('tag:yaml.org,2002:str', 247 data, 248 style='|') 249 return dumper.represent_scalar('tag:yaml.org,2002:str', data) 250 251 ConfigDumper.add_representer(str, str_presenter) 252 253 return ConfigDumper 254 255 256def main() -> None: 257 language_choices = sorted(scenario_config.LANGUAGES.keys()) 258 argp = argparse.ArgumentParser( 259 description='Generates load test configs from a template.', 260 fromfile_prefix_chars='@') 261 argp.add_argument('-l', 262 '--language', 263 action='append', 264 choices=language_choices, 265 required=True, 266 help='Language(s) to benchmark.', 267 dest='languages') 268 argp.add_argument('-t', 269 '--template', 270 type=str, 271 required=True, 272 help='LoadTest configuration yaml file template.') 273 argp.add_argument('-s', 274 '--substitution', 275 action='append', 276 default=[], 277 help='Template substitution(s), in the form key=value.', 278 dest='substitutions') 279 argp.add_argument('-p', 280 '--prefix', 281 default='', 282 type=str, 283 help='Test name prefix.') 284 argp.add_argument('-u', 285 '--uniquifier_element', 286 action='append', 287 default=[], 288 help='String element(s) to make the test name unique.', 289 dest='uniquifier_elements') 290 argp.add_argument( 291 '-d', 292 action='store_true', 293 help='Use creation date and time as an additional uniquifier element.') 294 argp.add_argument('-a', 295 '--annotation', 296 action='append', 297 default=[], 298 help='metadata.annotation(s), in the form key=value.', 299 dest='annotations') 300 argp.add_argument('-r', 301 '--regex', 302 default='.*', 303 type=str, 304 help='Regex to select scenarios to run.') 305 argp.add_argument( 306 '--category', 307 choices=['all', 'inproc', 'scalable', 'smoketest', 'sweep'], 308 default='all', 309 help='Select a category of tests to run.') 310 argp.add_argument( 311 '--allow_client_language', 312 action='append', 313 choices=language_choices, 314 default=[], 315 help='Allow cross-language scenarios with this client language.', 316 dest='allow_client_languages') 317 argp.add_argument( 318 '--allow_server_language', 319 action='append', 320 choices=language_choices, 321 default=[], 322 help='Allow cross-language scenarios with this server language.', 323 dest='allow_server_languages') 324 argp.add_argument('--runs_per_test', 325 default=1, 326 type=int, 327 help='Number of copies to generate for each test.') 328 argp.add_argument('-o', 329 '--output', 330 type=str, 331 help='Output file name. Output to stdout if not set.') 332 args = argp.parse_args() 333 334 substitutions = parse_key_value_args(args.substitutions) 335 336 uniquifier_elements = args.uniquifier_elements 337 if args.d: 338 uniquifier_elements.append(now_string()) 339 340 annotations = parse_key_value_args(args.annotations) 341 342 with open(args.template) as f: 343 base_config = yaml.safe_load( 344 string.Template(f.read()).substitute(substitutions)) 345 346 spec = base_config['spec'] 347 base_config_clients = spec['clients'] 348 del spec['clients'] 349 base_config_servers = spec['servers'] 350 del spec['servers'] 351 352 client_languages = [''] + args.allow_client_languages 353 server_languages = [''] + args.allow_server_languages 354 config_generators = [] 355 for l, cl, sl in itertools.product(args.languages, client_languages, 356 server_languages): 357 language_config = scenario_config_exporter.LanguageConfig( 358 category=args.category, 359 language=l, 360 client_language=cl, 361 server_language=sl) 362 config_generators.append( 363 gen_loadtest_configs(base_config, 364 base_config_clients, 365 base_config_servers, 366 args.regex, 367 language_config, 368 loadtest_name_prefix=args.prefix, 369 uniquifier_elements=uniquifier_elements, 370 annotations=annotations, 371 runs_per_test=args.runs_per_test)) 372 configs = (config for config in itertools.chain(*config_generators)) 373 374 with open(args.output, 'w') if args.output else sys.stdout as f: 375 yaml.dump_all(configs, 376 stream=f, 377 Dumper=config_dumper( 378 CONFIGURATION_FILE_HEADER_COMMENT.strip()), 379 default_flow_style=False) 380 381 382if __name__ == '__main__': 383 main() 384