@@ -52,34 +52,46 @@ def main(directory, additional_env_vars, save_name, compare_names, filter):
5252 benchmarks = [benchmark for benchmark in benchmarks if filter .search (benchmark .name ())]
5353
5454 for benchmark in benchmarks :
55- print (f"setting up { benchmark .name ()} ... " , end = '' , flush = True )
56- benchmark .setup ()
57- print ("complete." )
55+ try :
56+ print (f"setting up { benchmark .name ()} ... " , end = '' , flush = True )
57+ benchmark .setup ()
58+ print ("complete." )
59+ except Exception as e :
60+ if options .exit_on_failure :
61+ raise e
62+ else :
63+ print (f"failed: { e } " )
5864
5965 results = []
6066 for benchmark in benchmarks :
61- merged_env_vars = {** additional_env_vars }
62- iteration_results = []
63- for iter in range (options .iterations ):
64- print (f"running { benchmark .name ()} , iteration { iter } ... " , end = '' , flush = True )
65- bench_results = benchmark .run (merged_env_vars )
66- if bench_results is not None :
67- print (f"complete ({ bench_results .value } { benchmark .unit ()} )." )
68- iteration_results .append (bench_results )
67+ try :
68+ merged_env_vars = {** additional_env_vars }
69+ iteration_results = []
70+ for iter in range (options .iterations ):
71+ print (f"running { benchmark .name ()} , iteration { iter } ... " , end = '' , flush = True )
72+ bench_results = benchmark .run (merged_env_vars )
73+ if bench_results is not None :
74+ print (f"complete ({ bench_results .value } { benchmark .unit ()} )." )
75+ iteration_results .append (bench_results )
76+ else :
77+ print (f"did not finish." )
78+
79+ if len (iteration_results ) == 0 :
80+ continue
81+
82+ iteration_results .sort (key = lambda res : res .value )
83+ median_index = len (iteration_results ) // 2
84+ median_result = iteration_results [median_index ]
85+
86+ median_result .unit = benchmark .unit ()
87+ median_result .name = benchmark .name ()
88+
89+ results .append (median_result )
90+ except Exception as e :
91+ if options .exit_on_failure :
92+ raise e
6993 else :
70- print (f"did not finish." )
71-
72- if len (iteration_results ) == 0 :
73- continue
74-
75- iteration_results .sort (key = lambda res : res .value )
76- median_index = len (iteration_results ) // 2
77- median_result = iteration_results [median_index ]
78-
79- median_result .unit = benchmark .unit ()
80- median_result .name = benchmark .name ()
81-
82- results .append (median_result )
94+ print (f"failed: { e } " )
8395
8496 for benchmark in benchmarks :
8597 print (f"tearing down { benchmark .name ()} ... " , end = '' , flush = True )
@@ -126,6 +138,7 @@ def validate_and_parse_env_args(env_args):
126138 parser .add_argument ("--timeout" , type = int , help = 'Timeout for individual benchmarks in seconds.' , default = 600 )
127139 parser .add_argument ("--filter" , type = str , help = 'Regex pattern to filter benchmarks by name.' , default = None )
128140 parser .add_argument ("--verbose" , help = 'Print output of all the commands.' , action = "store_true" )
141+ parser .add_argument ("--exit_on_failure" , help = 'Exit on first failure.' , action = "store_true" )
129142
130143 args = parser .parse_args ()
131144 additional_env_vars = validate_and_parse_env_args (args .env )
@@ -137,6 +150,7 @@ def validate_and_parse_env_args(env_args):
137150 options .timeout = args .timeout
138151 options .ur_dir = args .ur_dir
139152 options .ur_adapter_name = args .ur_adapter_name
153+ options .exit_on_failure = args .exit_on_failure
140154
141155 benchmark_filter = re .compile (args .filter ) if args .filter else None
142156
0 commit comments