1import os
2import pyperf
3import tempfile
4
5
6def get_raw_values(filename, run_id):
7    bench = pyperf.Benchmark.load(filename)
8    run = bench.get_runs()[run_id]
9    inner_loops = run.get_inner_loops()
10    raw_values = [value * (loops * inner_loops)
11                  for loops, value in run.warmups]
12    total_loops = run.get_total_loops()
13    raw_values.extend(value * total_loops for value in run.values)
14    return (run, raw_values)
15
16
17class Replay(object):
18    def __init__(self, runner, filename):
19        self.runner = runner
20        self.args = runner.args
21        self.filename = filename
22        self.value_id = 0
23        self.init()
24
25    def init(self):
26        args = runner.args
27        self.run_id = self.args.first_run - 1
28        if args.worker:
29            self.read_session()
30
31            run, self.raw_values = get_raw_values(self.filename, self.run_id)
32            old_loops = args.loops
33            args.loops = run.get_loops()
34            if args.loops != old_loops:
35                print("Set loops to %s: value from the JSON" % args.loops)
36            # FIXME: handle inner_loops
37            self.run_id += 1
38            self.write_session()
39        else:
40            args.session_filename = tempfile.mktemp()
41            self.write_session()
42
43    def read_session(self):
44        filename = self.args.session_filename
45        if not filename:
46            return
47        with open(filename, "r") as fp:
48            line = fp.readline()
49        self.run_id = int(line.rstrip())
50
51    def write_session(self):
52        filename = self.args.session_filename
53        if not filename:
54            return
55        with open(filename, "w") as fp:
56            print(self.run_id, file=fp)
57            fp.flush()
58
59    def time_func(self, loops):
60        raw_value = self.raw_values[self.value_id]
61        self.value_id += 1
62        return raw_value
63
64
65def add_cmdline_args(cmd, args):
66    cmd.append(args.filename)
67    if args.session_filename:
68        cmd.extend(('--session-filename', args.session_filename))
69
70
71runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
72runner.argparser.add_argument('filename')
73runner.argparser.add_argument('--session-filename', default=None)
74runner.argparser.add_argument('--first-run', type=int, default=1)
75
76args = runner.parse_args()
77replay = Replay(runner, args.filename)
78runner.bench_time_func('bench', replay.time_func)
79if not args.worker:
80    os.unlink(args.session_filename)
81