diff options
Diffstat (limited to 'grpc/tools/run_tests/performance/loadtest_template.py')
| -rwxr-xr-x | grpc/tools/run_tests/performance/loadtest_template.py | 216 |
1 files changed, 216 insertions, 0 deletions
diff --git a/grpc/tools/run_tests/performance/loadtest_template.py b/grpc/tools/run_tests/performance/loadtest_template.py new file mode 100755 index 00000000..9684bb9a --- /dev/null +++ b/grpc/tools/run_tests/performance/loadtest_template.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +# Copyright 2021 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. + +# This script generates a load test configuration template from a collection of +# load test configurations. +# +# Configuration templates contain client and server configurations for multiple +# languages, and may contain template substitution keys. These templates are +# used to generate load test configurations by selecting clients and servers for +# the required languages. The source files for template generation may be load +# test configurations or load test configuration templates. Load test +# configuration generation is performed by loadtest_config.py. See documentation +# below: +# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md + +import argparse +import sys + +from typing import Any, Dict, Iterable, Mapping, Type + +import yaml + +import loadtest_config + +TEMPLATE_FILE_HEADER_COMMENT = """ +# Template generated from load test configurations by loadtest_template.py. +# +# Configuration templates contain client and server configurations for multiple +# languages, and may contain template substitution keys. These templates are +# used to generate load test configurations by selecting clients and servers for +# the required languages. The source files for template generation may be load +# test configurations or load test configuration templates. Load test +# configuration generation is performed by loadtest_config.py. See documentation +# below: +# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md +""" + + +def loadtest_template( + input_file_names: Iterable[str], + metadata: Mapping[str, Any], + inject_client_pool: bool, + inject_server_pool: bool, + inject_big_query_table: bool, + inject_timeout_seconds: bool, + inject_ttl_seconds: bool) -> Dict[str, Any]: # yapf: disable + """Generates the load test template.""" + clients = list() + servers = list() + spec = dict() + client_languages = set() + server_languages = set() + template = { + 'apiVersion': 'e2etest.grpc.io/v1', + 'kind': 'LoadTest', + 'metadata': metadata, + } + for input_file_name in input_file_names: + with open(input_file_name) as f: + input_config = yaml.safe_load(f.read()) + + if input_config.get('apiVersion') != template['apiVersion']: + raise ValueError('Unexpected api version in file {}: {}'.format( + input_file_name, input_config.get('apiVersion'))) + if input_config.get('kind') != template['kind']: + raise ValueError('Unexpected kind in file {}: {}'.format( + input_file_name, input_config.get('kind'))) + + for client in input_config['spec']['clients']: + if client['language'] in client_languages: + continue + if inject_client_pool: + client['pool'] = '${client_pool}' + clients.append(client) + client_languages.add(client['language']) + + for server in input_config['spec']['servers']: + if server['language'] in server_languages: + continue + if inject_server_pool: + server['pool'] = '${server_pool}' + servers.append(server) + server_languages.add(server['language']) + + input_spec = input_config['spec'] + del input_spec['clients'] + del input_spec['servers'] + del input_spec['scenariosJSON'] + spec.update(input_config['spec']) + + clients.sort(key=lambda x: x['language']) + servers.sort(key=lambda x: x['language']) + + spec.update({ + 'clients': clients, + 'servers': servers, + }) + + if inject_big_query_table: + if 'results' not in spec: + spec['results'] = dict() + spec['results']['bigQueryTable'] = '${big_query_table}' + if inject_timeout_seconds: + spec['timeoutSeconds'] = '${timeout_seconds}' + if inject_ttl_seconds: + spec['ttlSeconds'] = '${ttl_seconds}' + + template['spec'] = spec + + return template + + +def template_dumper(header_comment: str) -> Type[yaml.SafeDumper]: + """Returns a custom dumper to dump templates in the expected format.""" + + class TemplateDumper(yaml.SafeDumper): + + def expect_stream_start(self): + super().expect_stream_start() + if isinstance(self.event, yaml.StreamStartEvent): + self.write_indent() + self.write_indicator(header_comment, need_whitespace=False) + + def expect_block_sequence(self): + super().expect_block_sequence() + self.increase_indent() + + def expect_block_sequence_item(self, first=False): + if isinstance(self.event, yaml.SequenceEndEvent): + self.indent = self.indents.pop() + super().expect_block_sequence_item(first) + + return TemplateDumper + + +def main() -> None: + argp = argparse.ArgumentParser( + description='Creates a load test config generator template.', + fromfile_prefix_chars='@') + argp.add_argument('-i', + '--inputs', + action='extend', + nargs='+', + type=str, + help='Input files.') + argp.add_argument('-o', + '--output', + type=str, + help='Output file. Outputs to stdout if not set.') + argp.add_argument( + '--inject_client_pool', + action='store_true', + help='Set spec.client(s).pool values to \'${client_pool}\'.') + argp.add_argument( + '--inject_server_pool', + action='store_true', + help='Set spec.server(s).pool values to \'${server_pool}\'.') + argp.add_argument( + '--inject_big_query_table', + action='store_true', + help='Set spec.results.bigQueryTable to \'${big_query_table}\'.') + argp.add_argument('--inject_timeout_seconds', + action='store_true', + help='Set spec.timeoutSeconds to \'${timeout_seconds}\'.') + argp.add_argument('--inject_ttl_seconds', + action='store_true', + help='Set timeout ') + argp.add_argument('-n', + '--name', + default='', + type=str, + help='metadata.name.') + argp.add_argument('-a', + '--annotation', + action='append', + type=str, + help='metadata.annotation(s), in the form key=value.', + dest='annotations') + args = argp.parse_args() + + annotations = loadtest_config.parse_key_value_args(args.annotations) + + metadata = {'name': args.name} + if annotations: + metadata['annotations'] = annotations + + template = loadtest_template( + input_file_names=args.inputs, + metadata=metadata, + inject_client_pool=args.inject_client_pool, + inject_server_pool=args.inject_server_pool, + inject_big_query_table=args.inject_big_query_table, + inject_timeout_seconds=args.inject_timeout_seconds, + inject_ttl_seconds=args.inject_ttl_seconds) + + with open(args.output, 'w') if args.output else sys.stdout as f: + yaml.dump(template, + stream=f, + Dumper=template_dumper(TEMPLATE_FILE_HEADER_COMMENT.strip()), + default_flow_style=False) + + +if __name__ == '__main__': + main() |