benchmarks: added option '-o'/'--output' which controls where the benchmark...

benchmarks: added option '-o'/'--output' which controls where the benchmark results are printed (default is still stdout).

src/pipedream/benchmark/common.py

@@ -17,6 +17,7 @@ import sys
 import tempfile
 import time
 import typing as ty
+from typing.io import IO
 
 from pipedream.benchmark.types import *
 
@@ -32,7 +33,8 @@ from pipedream.utils import terminal
 def main(*, argv: ty.List[str] = None,
          make_perf_counters: ty.Callable[[ty.List[str]], 'Perf_Counter_Spec'],
          benchmark_kind: Benchmark_Kind):
-  _Benchmark_Runner().main(argv=argv, make_perf_counters=make_perf_counters, benchmark_kind=benchmark_kind)
+  _Benchmark_Runner().main(
+    argv=argv, make_perf_counters=make_perf_counters, benchmark_kind=benchmark_kind)
 
 
 def run_benchmarks(*,
@@ -78,6 +80,8 @@ class _Benchmark_Runner:
 
     perf_counters = make_perf_counters(args.run.extra_events)
 
+    output: IO[str] = args.output.output
+
     for results in run_benchmarks(arch                  = args.architecture,
                                   benchmarks            = benchmark_specs,
                                   perf_counters         = perf_counters,
@@ -88,10 +92,10 @@ class _Benchmark_Runner:
                                   tmp_dir               = args.output.tmp_dir,
                                   debug                 = args.debug,):
       self.info.clear()
-      sys.stdout.write('---\n')
-      yaml.dump(results, sys.stdout)
-      sys.stdout.write('...\n')
-      sys.stdout.flush()
+      output.write('---\n')
+      yaml.dump(results, output)
+      output.write('...\n')
+      output.flush()
 
   def run_benchmarks(self, *,
                      arch: ir.Architecture,
@@ -377,6 +381,16 @@ class _Benchmark_Runner:
 
     output = parser.add_argument_group('output options')
 
+    output_file = output.add_argument(
+      '-o', '--output',
+      help='''file to print benchmark results to''',
+    )
+    output_append = output.add_argument(
+      '-a', '--append',
+      help='''append to output file insted of overwriting it''',
+      action='store_true', default=False,
+    )
+
     tmp_dir = output.add_argument(
       '--tmp-dir',
       help='''directory for temporary files''',
@@ -438,6 +452,17 @@ class _Benchmark_Runner:
 
     ##
 
+    if not args.output:
+      args.output = sys.stdout
+    else:
+      mode = 'a' if args.append else 'w'
+      try:
+        args.output = argparse.FileType(mode)(args.output)
+      except argparse.ArgumentTypeError as e:
+        parser.error(str(e))
+
+    ##
+
     class Immutable_Bag(collections.abc.Mapping):
       def __init__(self, **kwargs):
         self.__dict__.update(kwargs)
@@ -470,7 +495,7 @@ class _Benchmark_Runner:
         kernels, tags, kernel_size, unroll_factor, different_instructions, dynamic_instructions, num_iterations,
         num_warmup_iterations, outlier_lo, outlier_hi, events),
       output=Immutable_Namespace(
-        tmp_dir, debug)
+        tmp_dir, debug, output_file)
     )
     return out