Improve benchmark script.

Author: hanke0

tools/runbenchmark.py

@@ -34,21 +34,37 @@ def humanize(t):
 
 
 class Benchmark:
-    def __init__(self, path):
-        self.path = os.path.abspath(path)
+    def __init__(self, dirs, error_exit=False):
+        self.files = []
+        for i in dirs:
+            self.files.extend(
+                glob.glob(os.path.join(i, "**/benchmark_*.py"), recursive=True)
+            )
 
-        self.files = glob.glob(os.path.join(path, "**/benchmark_*.py"), recursive=True)
         self.rv = {}
         self.finish = False
+        self.error_exit = error_exit
 
     def _exec_file(self, file):
         with open(file, "rt") as f:
             s = f.read()
         g = {"benchmark_setup": benchmark_setup}
-        exec(s, g)
+        try:
+            exec(s, g)
+        except Exception as e:
+            sys.stderr.write("Error load file ")
+            sys.stderr.write(file)
+            sys.stderr.write(": ")
+            sys.stderr.write(repr(e))
+            sys.stderr.write("\n")
+            sys.stderr.flush()
+            if self.error_exit:
+                sys.exit(1)
+            return None
         return g
 
-    def _is_benchmark(self, name, obj):
+    @staticmethod
+    def _is_benchmark(name, obj):
         if not inspect.isfunction(obj):
             return False
         if name == "benchmark_setup":
@@ -64,19 +80,36 @@ def _filter_benchmarks(self, g):
         return rv
 
     def start(self):
-        print("rootdir:", self.path)
-        print("Collect %d files" % len(self.files), flush=True)
-        for file in self.files:
+        print("Total %d files\n" % len(self.files), flush=True)
+        for idx, file in enumerate(self.files):
             g = self._exec_file(file)
+            if g is None:
+                continue
+
             benchmarks = self._filter_benchmarks(g)
-            print(
-                "\nFound {} benchmarks in file {}".format(len(benchmarks), file),
-                flush=True,
-            )
+            title = "Collect {} benchmarks in {}".format(len(benchmarks), file)
+            print(title, flush=True)
+            print("-" * len(title), flush=True)
             for name in benchmarks:
-                b_rv = self.run_str(name, g)
-                s = "{}::{}    {}".format(file, name, self._pretty_rv(*b_rv))
+                try:
+                    b_rv = self.run_str(name, g)
+                except Exception as e:
+                    sys.stderr.write("Error run benchmark '")
+                    sys.stderr.write(name)
+                    sys.stderr.write("' in file ")
+                    sys.stderr.write(file)
+                    sys.stderr.write(": ")
+                    sys.stderr.write(repr(e))
+                    sys.stderr.write("\n")
+                    sys.stderr.flush()
+                    if self.error_exit:
+                        sys.exit(1)
+                    else:
+                        continue
+
+                s = "* {}::{}    {}".format(file, name, self._pretty_rv(*b_rv))
                 print(s, flush=True)
+            print()
 
     def run_str(self, name, g):
         call = g[name]
@@ -88,21 +121,22 @@ def run_str(self, name, g):
 
         setup = ";".join(lines)
 
-        repeat = 10
-        total_cost = 8
+        repeat = 1
+        max_cost = 8
+        mask = 20
         cost = sum(timeit.repeat(callstr, setup=setup, globals=g, number=1, repeat=1))
-        if cost > total_cost:
-            loop = 3
+        if cost >= max_cost:
+            loop = 1
+        elif cost * mask >= max_cost:
+            loop = int(round(max_cost / cost))
+            loop += loop % 10
         else:
-            num = int(round(total_cost / cost))
+            num = int(round(max_cost / cost))
+            while num > mask * repeat and repeat < 10:
+                repeat += 1
+
             loop = num // repeat
-            if loop <= 3:
-                loop = 3
-            else:
-                count = 10
-                while count < loop:
-                    count *= 10
-                loop = count
+            loop += loop % 10
 
         costs = timeit.repeat(
             callstr, setup=setup, globals=g, number=loop, repeat=repeat
@@ -113,25 +147,34 @@ def run_str(self, name, g):
     @staticmethod
     def _pretty_rv(loop, repeat, costs):
         mean, std = mean_std(costs)
-        return "{} ± {}\t(each {:,} runs, {:,} loops)".format(
-            humanize(mean), humanize(std), repeat, loop
+        return "{} ± {}\t({:,} loops, each {:,} runs)".format(
+            humanize(mean), humanize(std), loop, repeat
         )
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("path", default=".")
-    if len(sys.argv) != 2:
-        path = "."
-    else:
-        path = sys.argv[1]
-        if path.startswith("-"):
-            parser.print_help()
-            sys.exit(1)
-
-    benchmark = Benchmark(path)
-
+    parser.add_argument("path", default=["."], nargs="+")
+    parser.add_argument("-p", "--python-path", default=["."], nargs="*")
+    parser.add_argument(
+        "-e",
+        "--error-exit",
+        default=False,
+        action="store_true",
+        help="Exit if error occurred",
+    )
+
+    args = parser.parse_args()
+    path = args.path
+    py_path = args.python_path
+    error_exit = args.error_exit
+    if py_path:
+        for i in py_path:
+            sys.path.insert(0, i)
+
+    benchmark = Benchmark(path, error_exit)
     benchmark.start()
+
     sys.exit(0)