• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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