#!/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 os import sys from typing import Any, Dict, Iterable, List, Mapping, Type import yaml sys.path.append(os.path.dirname(os.path.abspath(__file__))) 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 insert_worker(worker: Dict[str, Any], workers: List[Dict[str, Any]]) -> None: """Inserts client or server into a list, without inserting duplicates.""" def dump(w): return yaml.dump(w, Dumper=yaml.SafeDumper, default_flow_style=False) worker_str = dump(worker) if any((worker_str == dump(w) for w in workers)): return workers.append(worker) def uniquify_workers(workermap: Dict[str, List[Dict[str, Any]]]) -> None: """Name workers if there is more than one for the same map key.""" for workers in list(workermap.values()): if len(workers) <= 1: continue for i, worker in enumerate(workers): worker['name'] = str(i) def loadtest_template( input_file_names: Iterable[str], metadata: Mapping[str, Any], inject_client_pool: bool, inject_driver_image: bool, inject_driver_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.""" spec = dict() # type: Dict[str, Any] clientmap = dict() # Dict[str, List[Dict[str, Any]]] servermap = dict() # Dict[Str, List[Dict[str, Any]]] 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']: del client['name'] if inject_client_pool: client['pool'] = '${client_pool}' if client['language'] not in clientmap: clientmap[client['language']] = [] insert_worker(client, clientmap[client['language']]) for server in input_config['spec']['servers']: del server['name'] if inject_server_pool: server['pool'] = '${server_pool}' if server['language'] not in servermap: servermap[server['language']] = [] insert_worker(server, servermap[server['language']]) input_spec = input_config['spec'] del input_spec['clients'] del input_spec['servers'] del input_spec['scenariosJSON'] spec.update(input_config['spec']) uniquify_workers(clientmap) uniquify_workers(servermap) spec.update({ 'clients': sum((clientmap[language] for language in sorted(clientmap)), start=[]), 'servers': sum((servermap[language] for language in sorted(servermap)), start=[]), }) if 'driver' not in spec: spec['driver'] = {'language': 'cxx'} driver = spec['driver'] if 'name' in driver: del driver['name'] if inject_driver_image: if 'run' not in driver: driver['run'] = [{'name': 'main'}] driver['run'][0]['image'] = '${driver_image}' if inject_driver_pool: driver['pool'] = '${driver_pool}' if 'run' not in driver: if inject_driver_pool: raise ValueError('Cannot inject driver.pool: missing driver.run.') del spec['driver'] 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 str_presenter(dumper, data): if '\n' in data: return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|') return dumper.represent_scalar('tag:yaml.org,2002:str', data) TemplateDumper.add_representer(str, str_presenter) 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_driver_image', action='store_true', help='Set spec.driver(s).image values to \'${driver_image}\'.') argp.add_argument( '--inject_driver_pool', action='store_true', help='Set spec.driver(s).pool values to \'${driver_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_driver_image=args.inject_driver_image, inject_driver_pool=args.inject_driver_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()