Query CPU cost of tests and feed that into test runner

1 files changed, 16 insertions, 7 deletions

diff --git a/tools/run_tests/run_tests.py b/tools/run_tests/run_tests.py
index 4311e53839..6bb61c0b70 100755
--- a/tools/run_tests/run_tests.py
+++ b/tools/run_tests/run_tests.py
@@ -69,17 +69,20 @@ _POLLING_STRATEGIES = {
 }
 
 
-def get_flaky_tests(limit=None):
+BigQueryTestData = collections.namedtuple('BigQueryTestData', 'name flaky cpu')
+
+
+def get_bqtest_data(limit=None):
   import big_query_utils
 
   bq = big_query_utils.create_big_query()
   query = """
 SELECT
-  filtered_test_name,
+  filtered_test_name, SUM(result != 'PASSED' AND result != 'SKIPPED') > 0 as flaky, AVG(cpu_measured) as cpu
   FROM (
   SELECT
     REGEXP_REPLACE(test_name, r'/\d+', '') AS filtered_test_name,
-    result
+    result, cpu_measured
   FROM
     [grpc-testing:jenkins_test_results.aggregate_results]
   WHERE
@@ -89,15 +92,15 @@ SELECT
 GROUP BY
   filtered_test_name
 HAVING
-  SUM(result != 'PASSED' AND result != 'SKIPPED') > 0"""
+  flaky OR cpu > 0"""
   if limit:
     query += " limit {}".format(limit)
   query_job = big_query_utils.sync_query_job(bq, 'grpc-testing', query)
   page = bq.jobs().getQueryResults(
       pageToken=None,
       **query_job['jobReference']).execute(num_retries=3)
-  flake_names = [row['f'][0]['v'] for row in page['rows']]
-  return flake_names
+  test_data = [BigQueryTestData(row['f'][0]['v'], row['f'][1]['v'] == 'true', float(row['f'][2]['v'])) for row in page['rows']]
+  return test_data
 
 
 def platform_string():
@@ -141,6 +144,9 @@ class Config(object):
     if not flaky and shortname and shortname in flaky_tests:
       print('Setting %s to flaky' % shortname)
       flaky = True
+    if shortname in test_times:
+      print('Update CPU cost for %s: %f -> %f' % (shortname, cpu_cost, test_times[shortname]))
+      cpu_cost = test_times[shortname]
     return jobset.JobSpec(cmdline=self.tool_prefix + cmdline,
                           shortname=shortname,
                           environ=actual_environ,
@@ -1254,9 +1260,12 @@ argp.add_argument('--disable_auto_set_flakes', default=False, const=True, action
 args = argp.parse_args()
 
 flaky_tests = set()
+test_times = {}
 if not args.disable_auto_set_flakes:
   try:
-    flaky_tests = set(get_flaky_tests())
+    for test in get_bqtest_data():
+      if test.flaky: flaky_tests.add(test.name)
+      if test.cpu > 0: test_times[test.name] = test.cpu
   except:
     print("Unexpected error getting flaky tests:", sys.exc_info()[0])