1#!/usr/bin/env python2.7
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 bm_constants
19
20import argparse
21import subprocess
22import multiprocessing
23import random
24import itertools
25import sys
26import os
27
28sys.path.append(
29    os.path.join(
30        os.path.dirname(sys.argv[0]), '..', '..', '..', 'run_tests',
31        'python_utils'))
32import jobset
33
34
35def _args():
36    argp = argparse.ArgumentParser(description='Runs microbenchmarks')
37    argp.add_argument(
38        '-b',
39        '--benchmarks',
40        nargs='+',
41        choices=bm_constants._AVAILABLE_BENCHMARK_TESTS,
42        default=bm_constants._AVAILABLE_BENCHMARK_TESTS,
43        help='Benchmarks to run')
44    argp.add_argument(
45        '-j',
46        '--jobs',
47        type=int,
48        default=multiprocessing.cpu_count(),
49        help='Number of CPUs to use')
50    argp.add_argument(
51        '-n',
52        '--name',
53        type=str,
54        help=
55        'Unique name of the build to run. Needs to match the handle passed to bm_build.py'
56    )
57    argp.add_argument(
58        '-r',
59        '--regex',
60        type=str,
61        default="",
62        help='Regex to filter benchmarks run')
63    argp.add_argument(
64        '-l',
65        '--loops',
66        type=int,
67        default=20,
68        help=
69        'Number of times to loops the benchmarks. More loops cuts down on noise'
70    )
71    argp.add_argument('--counters', dest='counters', action='store_true')
72    argp.add_argument('--no-counters', dest='counters', action='store_false')
73    argp.set_defaults(counters=True)
74    args = argp.parse_args()
75    assert args.name
76    if args.loops < 3:
77        print "WARNING: This run will likely be noisy. Increase loops to at least 3."
78    return args
79
80
81def _collect_bm_data(bm, cfg, name, regex, idx, loops):
82    jobs_list = []
83    for line in subprocess.check_output([
84            'bm_diff_%s/%s/%s' % (name, cfg, bm), '--benchmark_list_tests',
85            '--benchmark_filter=%s' % regex
86    ]).splitlines():
87        stripped_line = line.strip().replace("/",
88                                             "_").replace("<", "_").replace(
89                                                 ">", "_").replace(", ", "_")
90        cmd = [
91            'bm_diff_%s/%s/%s' % (name, cfg, bm),
92            '--benchmark_filter=^%s$' % line,
93            '--benchmark_out=%s.%s.%s.%s.%d.json' % (bm, stripped_line, cfg,
94                                                     name, idx),
95            '--benchmark_out_format=json',
96        ]
97        jobs_list.append(
98            jobset.JobSpec(
99                cmd,
100                shortname='%s %s %s %s %d/%d' % (bm, line, cfg, name, idx + 1,
101                                                 loops),
102                verbose_success=True,
103                cpu_cost=2,
104                timeout_seconds=60 * 60))  # one hour
105    return jobs_list
106
107
108def create_jobs(name, benchmarks, loops, regex, counters):
109    jobs_list = []
110    for loop in range(0, loops):
111        for bm in benchmarks:
112            jobs_list += _collect_bm_data(bm, 'opt', name, regex, loop, loops)
113            if counters:
114                jobs_list += _collect_bm_data(bm, 'counters', name, regex, loop,
115                                              loops)
116    random.shuffle(jobs_list, random.SystemRandom().random)
117    return jobs_list
118
119
120if __name__ == '__main__':
121    args = _args()
122    jobs_list = create_jobs(args.name, args.benchmarks, args.loops, args.regex,
123                            args.counters)
124    jobset.run(jobs_list, maxjobs=args.jobs)
125