Add --exit-on-exception to exit process on any exception (meta-pytorch#460)

Author: yf225

tritonbench/utils/parser.py

@@ -154,6 +154,12 @@ def get_parser(args=None):
         "--keep-going",
         action="store_true",
     )
+    parser.add_argument(
+        "--exit-on-exception",
+        action="store_true",
+        default=False,
+        help="Immediately terminate the process if any operator run raises an exception.",
+    )
     parser.add_argument(
         "--input-id",
         type=int,

tritonbench/utils/triton_op.py

@@ -994,6 +994,8 @@ def _reduce_benchmarks(acc, bm_name: str):
                 "Caught exception, terminating early with partial results",
                 exc_info=True,
             )
+            if self.tb_args.exit_on_exception:
+                os._exit(1)
             raise
         finally:
             self.output = BenchmarkOperatorResult(