-
Notifications
You must be signed in to change notification settings - Fork 5
/
test.py
executable file
·96 lines (84 loc) · 2.92 KB
/
test.py
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
#!/usr/bin/python3
import argparse
import os
import subprocess
import sys
from build import print_header, STEP_IMPLEMENTATIONS
from bench import CPU_BIND_CMD
INPUT_SIZES = [1, 2, 31, 32, 33, 99, 100, 101, 120, 150, 175, 199, 201, 501, 999]
def run(cmd, num_threads):
env = None
if num_threads:
env = {
"OMP_NUM_THREADS": str(num_threads),
"RUST_BACKTRACE": "full",
"RAYON_NUM_THREADS": str(num_threads)
}
cmd = CPU_BIND_CMD + ' ' + cmd
result = subprocess.run(
cmd.split(' '),
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
check=True,
env=env
)
return result.stdout.decode("utf-8")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--build_dir", "-b",
type=str,
default=os.path.join(os.path.curdir, "build"),
help="Path to the benchmark binaries, if not the default from build.py.")
parser.add_argument("--threads", "-t",
type=int,
default=1,
help="Value for environment variables controlling number of threads")
parser.add_argument("--input_size", "-n",
type=int,
help="Run tests with this value for n")
parser.add_argument("--iterations", "-c",
type=int,
default=10,
help="Run each test this many times")
parser.add_argument("--implementation", "-i",
type=str,
help="Filter implementations by prefix, e.g '-i v0' runs only v0_baseline.")
parser.add_argument("--no-cpp",
action='store_true')
parser.add_argument("--no-rust",
action='store_true')
parser.add_argument("--verbose", "-v",
action='store_true',
default=False)
args = parser.parse_args()
build_dir = os.path.abspath(os.path.join(args.build_dir, "bin"))
impl_filter = args.implementation
if args.input_size:
INPUT_SIZES = [args.input_size]
all_ok = True
langs = []
if not args.no_cpp:
langs.append("cpp")
if not args.no_rust:
langs.append("rust")
for lang in langs:
for step_impl in STEP_IMPLEMENTATIONS:
if impl_filter and not step_impl.startswith(impl_filter):
continue
print_header(lang + ' ' + step_impl + ' ...', end=' ')
test_cmd = os.path.join(build_dir, step_impl + "_" + lang)
failed = False
for input_size in INPUT_SIZES:
test_args = 'test {} {}'.format(input_size, args.iterations)
cmd = test_cmd + ' ' + test_args
output = run(cmd, args.threads)
failed = "ERROR" in output
if args.verbose:
print(output)
if all_ok and failed:
all_ok = False
print("! fail" if failed else "ok")
if not all_ok:
print()
print_header("ERROR: at least one test failed")
sys.exit(1)