11#!/usr/bin/env python3
22# -*- coding: utf-8 -*-
33#
4- # SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
4+ # SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
55# SPDX-License-Identifier: Apache-2.0
66#
77# Licensed under the Apache License, Version 2.0 (the "License");
@@ -57,7 +57,7 @@ def __init__(self, spark_session: SparkSession, query_name) -> None:
5757 'query' : query_name ,
5858 }
5959
60- def report_on (self , fn : Callable , * args ):
60+ def report_on (self , fn : Callable , warmup_iterations = 0 , iterations = 1 , * args ):
6161 """Record a function for its running environment, running status etc. and exclude sentive
6262 information like tokens, secret and password Generate summary in dict format for it.
6363
@@ -84,28 +84,41 @@ def report_on(self, fn: Callable, *args):
8484 if listener is not None :
8585 print ("TaskFailureListener is registered." )
8686 try :
87- start_time = int (time .time () * 1000 )
88- fn (* args )
89- end_time = int (time .time () * 1000 )
90- if listener and len (listener .failures ) != 0 :
91- self .summary ['queryStatus' ].append ("CompletedWithTaskFailures" )
92- else :
93- self .summary ['queryStatus' ].append ("Completed" )
87+ # warmup
88+ for i in range (0 , warmup_iterations ):
89+ fn (* args )
9490 except Exception as e :
95- # print the exception to ease debugging
96- print ('ERROR BEGIN' )
91+ print ('ERROR WHILE WARMUP BEGIN' )
9792 print (e )
9893 traceback .print_tb (e .__traceback__ )
99- print ('ERROR END' )
100- end_time = int (time .time () * 1000 )
101- self .summary ['queryStatus' ].append ("Failed" )
102- self .summary ['exceptions' ].append (str (e ))
103- finally :
104- self .summary ['startTime' ] = start_time
105- self .summary ['queryTimes' ].append (end_time - start_time )
106- if listener is not None :
107- listener .unregister ()
108- return self .summary
94+ print ('ERROR WHILE WARMUP END' )
95+
96+ start_time = int (time .time () * 1000 )
97+ self .summary ['startTime' ] = start_time
98+ # run the query
99+ for i in range (0 , iterations ):
100+ try :
101+ start_time = int (time .time () * 1000 )
102+ fn (* args )
103+ end_time = int (time .time () * 1000 )
104+ if listener and len (listener .failures ) != 0 :
105+ self .summary ['queryStatus' ].append ("CompletedWithTaskFailures" )
106+ else :
107+ self .summary ['queryStatus' ].append ("Completed" )
108+ except Exception as e :
109+ # print the exception to ease debugging
110+ print ('ERROR BEGIN' )
111+ print (e )
112+ traceback .print_tb (e .__traceback__ )
113+ print ('ERROR END' )
114+ end_time = int (time .time () * 1000 )
115+ self .summary ['queryStatus' ].append ("Failed" )
116+ self .summary ['exceptions' ].append (str (e ))
117+ finally :
118+ self .summary ['queryTimes' ].append (end_time - start_time )
119+ if listener is not None :
120+ listener .unregister ()
121+ return self .summary
109122
110123 def write_summary (self , prefix = "" ):
111124 """_summary_
0 commit comments