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# This script generates a load test configuration template from a collection of 17# load test configurations. 18# 19# Configuration templates contain client and server configurations for multiple 20# languages, and may contain template substitution keys. These templates are 21# used to generate load test configurations by selecting clients and servers for 22# the required languages. The source files for template generation may be load 23# test configurations or load test configuration templates. Load test 24# configuration generation is performed by loadtest_config.py. See documentation 25# below: 26# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md 27 28import argparse 29import sys 30 31from typing import Any, Dict, Iterable, Mapping, Type 32 33import yaml 34 35import loadtest_config 36 37TEMPLATE_FILE_HEADER_COMMENT = """ 38# Template generated from load test configurations by loadtest_template.py. 39# 40# Configuration templates contain client and server configurations for multiple 41# languages, and may contain template substitution keys. These templates are 42# used to generate load test configurations by selecting clients and servers for 43# the required languages. The source files for template generation may be load 44# test configurations or load test configuration templates. Load test 45# configuration generation is performed by loadtest_config.py. See documentation 46# below: 47# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md 48""" 49 50 51def loadtest_template( 52 input_file_names: Iterable[str], 53 metadata: Mapping[str, Any], 54 inject_client_pool: bool, 55 inject_server_pool: bool, 56 inject_big_query_table: bool, 57 inject_timeout_seconds: bool, 58 inject_ttl_seconds: bool) -> Dict[str, Any]: # yapf: disable 59 """Generates the load test template.""" 60 clients = list() 61 servers = list() 62 spec = dict() 63 client_languages = set() 64 server_languages = set() 65 template = { 66 'apiVersion': 'e2etest.grpc.io/v1', 67 'kind': 'LoadTest', 68 'metadata': metadata, 69 } 70 for input_file_name in input_file_names: 71 with open(input_file_name) as f: 72 input_config = yaml.safe_load(f.read()) 73 74 if input_config.get('apiVersion') != template['apiVersion']: 75 raise ValueError('Unexpected api version in file {}: {}'.format( 76 input_file_name, input_config.get('apiVersion'))) 77 if input_config.get('kind') != template['kind']: 78 raise ValueError('Unexpected kind in file {}: {}'.format( 79 input_file_name, input_config.get('kind'))) 80 81 for client in input_config['spec']['clients']: 82 if client['language'] in client_languages: 83 continue 84 if inject_client_pool: 85 client['pool'] = '${client_pool}' 86 clients.append(client) 87 client_languages.add(client['language']) 88 89 for server in input_config['spec']['servers']: 90 if server['language'] in server_languages: 91 continue 92 if inject_server_pool: 93 server['pool'] = '${server_pool}' 94 servers.append(server) 95 server_languages.add(server['language']) 96 97 input_spec = input_config['spec'] 98 del input_spec['clients'] 99 del input_spec['servers'] 100 del input_spec['scenariosJSON'] 101 spec.update(input_config['spec']) 102 103 clients.sort(key=lambda x: x['language']) 104 servers.sort(key=lambda x: x['language']) 105 106 spec.update({ 107 'clients': clients, 108 'servers': servers, 109 }) 110 111 if inject_big_query_table: 112 if 'results' not in spec: 113 spec['results'] = dict() 114 spec['results']['bigQueryTable'] = '${big_query_table}' 115 if inject_timeout_seconds: 116 spec['timeoutSeconds'] = '${timeout_seconds}' 117 if inject_ttl_seconds: 118 spec['ttlSeconds'] = '${ttl_seconds}' 119 120 template['spec'] = spec 121 122 return template 123 124 125def template_dumper(header_comment: str) -> Type[yaml.SafeDumper]: 126 """Returns a custom dumper to dump templates in the expected format.""" 127 128 class TemplateDumper(yaml.SafeDumper): 129 130 def expect_stream_start(self): 131 super().expect_stream_start() 132 if isinstance(self.event, yaml.StreamStartEvent): 133 self.write_indent() 134 self.write_indicator(header_comment, need_whitespace=False) 135 136 def expect_block_sequence(self): 137 super().expect_block_sequence() 138 self.increase_indent() 139 140 def expect_block_sequence_item(self, first=False): 141 if isinstance(self.event, yaml.SequenceEndEvent): 142 self.indent = self.indents.pop() 143 super().expect_block_sequence_item(first) 144 145 return TemplateDumper 146 147 148def main() -> None: 149 argp = argparse.ArgumentParser( 150 description='Creates a load test config generator template.', 151 fromfile_prefix_chars='@') 152 argp.add_argument('-i', 153 '--inputs', 154 action='extend', 155 nargs='+', 156 type=str, 157 help='Input files.') 158 argp.add_argument('-o', 159 '--output', 160 type=str, 161 help='Output file. Outputs to stdout if not set.') 162 argp.add_argument( 163 '--inject_client_pool', 164 action='store_true', 165 help='Set spec.client(s).pool values to \'${client_pool}\'.') 166 argp.add_argument( 167 '--inject_server_pool', 168 action='store_true', 169 help='Set spec.server(s).pool values to \'${server_pool}\'.') 170 argp.add_argument( 171 '--inject_big_query_table', 172 action='store_true', 173 help='Set spec.results.bigQueryTable to \'${big_query_table}\'.') 174 argp.add_argument('--inject_timeout_seconds', 175 action='store_true', 176 help='Set spec.timeoutSeconds to \'${timeout_seconds}\'.') 177 argp.add_argument('--inject_ttl_seconds', 178 action='store_true', 179 help='Set timeout ') 180 argp.add_argument('-n', 181 '--name', 182 default='', 183 type=str, 184 help='metadata.name.') 185 argp.add_argument('-a', 186 '--annotation', 187 action='append', 188 type=str, 189 help='metadata.annotation(s), in the form key=value.', 190 dest='annotations') 191 args = argp.parse_args() 192 193 annotations = loadtest_config.parse_key_value_args(args.annotations) 194 195 metadata = {'name': args.name} 196 if annotations: 197 metadata['annotations'] = annotations 198 199 template = loadtest_template( 200 input_file_names=args.inputs, 201 metadata=metadata, 202 inject_client_pool=args.inject_client_pool, 203 inject_server_pool=args.inject_server_pool, 204 inject_big_query_table=args.inject_big_query_table, 205 inject_timeout_seconds=args.inject_timeout_seconds, 206 inject_ttl_seconds=args.inject_ttl_seconds) 207 208 with open(args.output, 'w') if args.output else sys.stdout as f: 209 yaml.dump(template, 210 stream=f, 211 Dumper=template_dumper(TEMPLATE_FILE_HEADER_COMMENT.strip()), 212 default_flow_style=False) 213 214 215if __name__ == '__main__': 216 main() 217