summaryrefslogtreecommitdiff
path: root/nightbench.py
blob: ffaccf9bfb64e834026bd61939a8967f85e33ef4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
#!/usr/bin/env python3
import argparse
import datetime
import math
import shlex
import signal
import subprocess
import sys
import time

def run_command(args, line):
    try:
        if args.split:
            subprocess.check_call(args.command + shlex.split(line))
        else:
            subprocess.check_call(args.command + [line])
    except subprocess.CalledProcessError as e:
        print(f"Command exited with code {e.returncode}, stopping nightbench", file=sys.stderr)
        sys.exit(1)

def parse_iv_endpoint(arg):
    try:
        tm = time.strptime(arg, "%Y-%m-%d %H:%M:%S")
        return datetime.datetime(tm.tm_year, tm.tm_mon, tm.tm_mday, tm.tm_hour, tm.tm_min, tm.tm_sec)
    except ValueError as e:
        pass

    try:
        tm = time.strptime(arg, "%H:%M:%S")
        return datetime.time(tm.tm_hour, tm.tm_min, tm.tm_sec)
    except ValueError as e:
        pass

    print(f"Failed to parse {arg!r} as a (date)time", file=sys.stderr)
    sys.exit(1)

def parse_interval(arg1, arg2):
    t1 = parse_iv_endpoint(arg1)
    t2 = parse_iv_endpoint(arg2)
    if isinstance(t1, datetime.datetime) and isinstance(t2, datetime.datetime):
        return (t1, t2)
    if isinstance(t1, datetime.time) and isinstance(t2, datetime.time):
        return (t1, t2)
    print(f"Interval has inconsistent endpoints: {arg1!r} {arg2!r}", file=sys.stderr)
    sys.exit(1)

# Returns (time to wait until allowed, remaining time allowed from now)
def time_until(ivs):
    if len(ivs) == 0: return 0.0, math.inf
    now = datetime.datetime.now()
    nowts = now.timestamp()
    nowtime = now.time()
    tomorrowdate = datetime.date.fromordinal(now.date().toordinal() + 1)
    minwait = math.inf
    time_remaining = 0.0
    for start, end in ivs:
        if isinstance(start, datetime.datetime):
            startts, endts = start.timestamp(), end.timestamp()
            if startts <= nowts < endts:
                minwait = 0.0
                time_remaining = max(time_remaining, endts - nowts)
            elif nowts < startts: minwait = min(minwait, startts - nowts)
        elif isinstance(start, datetime.time):
            if start <= now.time() < end or (end < start and (start <= now.time() or now.time() < end)):
                minwait = 0.0
                endts = datetime.datetime(now.year, now.month, now.day, end.hour, end.minute, end.second).timestamp()
                if endts <= nowts:
                    endts = datetime.datetime(
                            tomorrowdate.year, tomorrowdate.month, tomorrowdate.day,
                            end.hour, end.minute, end.second).timestamp()
                time_remaining = max(time_remaining, endts - nowts)
                continue
            startts = datetime.datetime(now.year, now.month, now.day, start.hour, start.minute, start.second).timestamp()
            if nowts < startts:
                minwait = min(minwait, startts - nowts)
                continue
            startts = datetime.datetime(
                    tomorrowdate.year, tomorrowdate.month, tomorrowdate.day,
                    start.hour, start.minute, start.second).timestamp()
            assert nowts < startts
            minwait = min(minwait, startts - nowts)
    return minwait, time_remaining

class Estimator:
    def __init__(self):
        self._sum = 0.0
        self._sumsq = 0.0
        self._n = 0

    def observe(self, time):
        self._sum += time
        self._sumsq += time * time
        self._n += 1

    def status(self):
        n = self._n
        if n < 4: return None, None
        # stddev = sqrt (1/(n-1) sum_i (xi - mu)^2)
        #        = sqrt (1/(n-1) (sum_i xi^2 - 1/n (sum_i xi)^2))
        #        = sqrt (1/(n-1) sum_i xi^2 - 1/(n(n-1)) (sum_i xi)^2)
        stddev = math.sqrt(self._sumsq / (n-1) - 1/(n*(n-1)) * self._sum ** 2)
        mu = self._sum / n
        # If variance is not outrageously big, try to ensure that even higher
        # values end up inside the valid interval
        if stddev < mu: return mu + stddev, "μ+σ"
        else: return mu, "μ"

    def estimate(self):
        return self.status()[0]


def main():
    help_description = """
Run your benchmarks with time constraints. Reads lines from stdin; the given
command is executed once for each line, with the line passed as a single
command line argument (but see -s). Nightbench attempts to ensure the command
runs only during the allowed time intervals by estimating command duration and
sleeping during disallowed intervals.

If no intervals are given, all times are assumed to be allowed.
"""

    help_epilog = """
The 'TIME' in -t is a time, possibly with a date. Specifically, a TIME must be
in one of these formats: (both ends must have the same format)
1. "HH:MM:SS"
2. "YYYY-mm-dd HH:MM:SS"
where format 1 is considered to be that time occurring every day. Times are in
the system local time and are in 24-hour format.
"""

    parser = argparse.ArgumentParser(
                prog='nightbench',
                description=help_description,
                epilog=help_epilog,
                formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("command", nargs='+',
                        help="The command to execute (directly, not via shell)")
    parser.add_argument("-t", action="append", metavar="TIME", nargs=2,
                        help="Allow runs in this interval")
    parser.add_argument("-s", "--split", action="store_true",
                        help="Shell-split lines into command-line arguments")
    args = parser.parse_args()
    #  print(args)

    intervals = [parse_interval(*iv) for iv in ([] if args.t is None else args.t)]
    #  print(intervals)

    estimator = Estimator()
    try:
        while True:
            # 1. Ascertain there's work
            line = sys.stdin.readline()
            if line == "": break  # EOF
            if line[-1] == "\n": line = line[:-1]

            # 2. Wait until we get to run
            while True:
                waittime, time_remaining = time_until(intervals)
                # If this is not the first run we get to do this time (waittime == 0)
                # and one more run won't fit, skip until the next allowed interval
                if waittime == 0.0 and \
                        (est := estimator.estimate()) is not None and \
                        time_remaining < est:
                    print(f"# nightbench: predicting a run won't fit in {time_remaining} seconds, sleeping {time_remaining} seconds", file=sys.stderr)
                    time.sleep(time_remaining + 0.1)
                    continue
                if waittime > 0.0:
                    print(f"# nightbench: sleeping {waittime} seconds", file=sys.stderr)
                    time.sleep(waittime)
                    if time_until(intervals)[0] > 0.1: continue  # ??
                break

            # 3. Do a run
            tm1 = time.clock_gettime(time.CLOCK_MONOTONIC)
            run_command(args, line)
            tm2 = time.clock_gettime(time.CLOCK_MONOTONIC)
            estimator.observe(tm2 - tm1)

            # 4. Report
            extrainfo = ""
            if (est := estimator.status())[0] is not None:
                extrainfo = f" (estimate {est[1]} = {est[0]})"
            print(f"# nightbench: run took {tm2 - tm1} seconds{extrainfo}", file=sys.stderr)
    except KeyboardInterrupt as e:
        print("Interrupted.", file=sys.stderr)
        sys.exit(128 + signal.SIGINT)


if __name__ == "__main__":
    main()