1#!/usr/bin/env python3
2#
3# Copyright 2017 gRPC authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16""" Python utility to run opt and counters benchmarks and save json output """
17
18import argparse
19import itertools
20import multiprocessing
21import os
22import random
23import subprocess
24import sys
25
26import bm_constants
27import jobset
28
29sys.path.append(
30    os.path.join(os.path.dirname(sys.argv[0]), '..', '..', '..', 'run_tests',
31                 'python_utils'))
32
33
34def _args():
35    argp = argparse.ArgumentParser(description='Runs microbenchmarks')
36    argp.add_argument('-b',
37                      '--benchmarks',
38                      nargs='+',
39                      choices=bm_constants._AVAILABLE_BENCHMARK_TESTS,
40                      default=bm_constants._AVAILABLE_BENCHMARK_TESTS,
41                      help='Benchmarks to run')
42    argp.add_argument('-j',
43                      '--jobs',
44                      type=int,
45                      default=multiprocessing.cpu_count(),
46                      help='Number of CPUs to use')
47    argp.add_argument(
48        '-n',
49        '--name',
50        type=str,
51        help=
52        'Unique name of the build to run. Needs to match the handle passed to bm_build.py'
53    )
54    argp.add_argument('-r',
55                      '--regex',
56                      type=str,
57                      default="",
58                      help='Regex to filter benchmarks run')
59    argp.add_argument(
60        '-l',
61        '--loops',
62        type=int,
63        default=20,
64        help=
65        'Number of times to loops the benchmarks. More loops cuts down on noise'
66    )
67    argp.add_argument('--counters', dest='counters', action='store_true')
68    argp.add_argument('--no-counters', dest='counters', action='store_false')
69    argp.set_defaults(counters=True)
70    args = argp.parse_args()
71    assert args.name
72    if args.loops < 3:
73        print("WARNING: This run will likely be noisy. Increase loops to at "
74              "least 3.")
75    return args
76
77
78def _collect_bm_data(bm, cfg, name, regex, idx, loops):
79    jobs_list = []
80    for line in subprocess.check_output([
81            'bm_diff_%s/%s/%s' % (name, cfg, bm), '--benchmark_list_tests',
82            '--benchmark_filter=%s' % regex
83    ]).splitlines():
84        line = line.decode('UTF-8')
85        stripped_line = line.strip().replace("/",
86                                             "_").replace("<", "_").replace(
87                                                 ">", "_").replace(", ", "_")
88        cmd = [
89            'bm_diff_%s/%s/%s' % (name, cfg, bm),
90            '--benchmark_filter=^%s$' % line,
91            '--benchmark_out=%s.%s.%s.%s.%d.json' %
92            (bm, stripped_line, cfg, name, idx),
93            '--benchmark_out_format=json',
94        ]
95        jobs_list.append(
96            jobset.JobSpec(cmd,
97                           shortname='%s %s %s %s %d/%d' %
98                           (bm, line, cfg, name, idx + 1, loops),
99                           verbose_success=True,
100                           cpu_cost=2,
101                           timeout_seconds=60 * 60))  # one hour
102    return jobs_list
103
104
105def create_jobs(name, benchmarks, loops, regex, counters):
106    jobs_list = []
107    for loop in range(0, loops):
108        for bm in benchmarks:
109            jobs_list += _collect_bm_data(bm, 'opt', name, regex, loop, loops)
110            if counters:
111                jobs_list += _collect_bm_data(bm, 'counters', name, regex, loop,
112                                              loops)
113    random.shuffle(jobs_list, random.SystemRandom().random)
114    return jobs_list
115
116
117if __name__ == '__main__':
118    args = _args()
119    jobs_list = create_jobs(args.name, args.benchmarks, args.loops, args.regex,
120                            args.counters)
121    jobset.run(jobs_list, maxjobs=args.jobs)
122