Add --inner_jobs args to run_tests_matrix.py

1 files changed, 51 insertions, 33 deletions

diff --git a/tools/run_tests/run_tests_matrix.py b/tools/run_tests/run_tests_matrix.py
index 21d3dd4a0b..af24d75236 100755
--- a/tools/run_tests/run_tests_matrix.py
+++ b/tools/run_tests/run_tests_matrix.py
@@ -46,16 +46,16 @@ os.chdir(_ROOT)
 _RUNTESTS_TIMEOUT = 2*60*60
 
 # Number of jobs assigned to each run_tests.py instance
-_INNER_JOBS = 2
+_DEFAULT_INNER_JOBS = 2
 
 
-def _docker_jobspec(name, runtests_args=[]):
+def _docker_jobspec(name, runtests_args=[], inner_jobs=_DEFAULT_INNER_JOBS):
   """Run a single instance of run_tests.py in a docker container"""
   test_job = jobset.JobSpec(
           cmdline=['python', 'tools/run_tests/run_tests.py',
                    '--use_docker',
                    '-t',
-                   '-j', str(_INNER_JOBS),
+                   '-j', str(inner_jobs),
                    '-x', 'report_%s.xml' % name,
                    '--report_suite_name', '%s' % name] + runtests_args,
           shortname='run_tests_%s' % name,
@@ -63,7 +63,7 @@ def _docker_jobspec(name, runtests_args=[]):
   return test_job
 
 
-def _workspace_jobspec(name, runtests_args=[], workspace_name=None):
+def _workspace_jobspec(name, runtests_args=[], workspace_name=None, inner_jobs=_DEFAULT_INNER_JOBS):
   """Run a single instance of run_tests.py in a separate workspace"""
   if not workspace_name:
     workspace_name = 'workspace_%s' % name
@@ -71,7 +71,7 @@ def _workspace_jobspec(name, runtests_args=[], workspace_name=None):
   test_job = jobset.JobSpec(
           cmdline=['tools/run_tests/run_tests_in_workspace.sh',
                    '-t',
-                   '-j', str(_INNER_JOBS),
+                   '-j', str(inner_jobs),
                    '-x', '../report_%s.xml' % name,
                    '--report_suite_name', '%s' % name] + runtests_args,
           environ=env,
@@ -82,7 +82,8 @@ def _workspace_jobspec(name, runtests_args=[], workspace_name=None):
 
 def _generate_jobs(languages, configs, platforms,
                   arch=None, compiler=None,
-                  labels=[], extra_args=[]):
+                  labels=[], extra_args=[],
+                  inner_jobs=_DEFAULT_INNER_JOBS):
   result = []
   for language in languages:
     for platform in platforms:
@@ -97,60 +98,66 @@ def _generate_jobs(languages, configs, platforms,
 
         runtests_args += extra_args
         if platform == 'linux':
-          job = _docker_jobspec(name=name, runtests_args=runtests_args)
+          job = _docker_jobspec(name=name, runtests_args=runtests_args, inner_jobs=inner_jobs)
         else:
-          job = _workspace_jobspec(name=name, runtests_args=runtests_args)
+          job = _workspace_jobspec(name=name, runtests_args=runtests_args, inner_jobs=inner_jobs)
 
         job.labels = [platform, config, language] + labels
         result.append(job)
   return result
 
 
-def _create_test_jobs(extra_args=[]):
+def _create_test_jobs(extra_args=[], inner_jobs=_DEFAULT_INNER_JOBS):
   test_jobs = []
   # supported on linux only
   test_jobs += _generate_jobs(languages=['sanity', 'php7'],
                              configs=['dbg', 'opt'],
                              platforms=['linux'],
                              labels=['basictests'],
-                             extra_args=extra_args)
-  
+                             extra_args=extra_args,
+                             inner_jobs=inner_jobs)
+
   # supported on all platforms.
   test_jobs += _generate_jobs(languages=['c', 'csharp', 'node', 'python'],
                              configs=['dbg', 'opt'],
                              platforms=['linux', 'macos', 'windows'],
                              labels=['basictests'],
-                             extra_args=extra_args)
-  
+                             extra_args=extra_args,
+                             inner_jobs=inner_jobs)
+
   # supported on linux and mac.
   test_jobs += _generate_jobs(languages=['c++', 'ruby', 'php'],
                               configs=['dbg', 'opt'],
                               platforms=['linux', 'macos'],
                               labels=['basictests'],
-                              extra_args=extra_args)
-  
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
+
   # supported on mac only.
   test_jobs += _generate_jobs(languages=['objc'],
                               configs=['dbg', 'opt'],
                               platforms=['macos'],
                               labels=['basictests'],
-                              extra_args=extra_args)
-  
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
+
   # sanitizers
   test_jobs += _generate_jobs(languages=['c'],
                               configs=['msan', 'asan', 'tsan'],
                               platforms=['linux'],
                               labels=['sanitizers'],
-                              extra_args=extra_args)
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
   test_jobs += _generate_jobs(languages=['c++'],
                               configs=['asan', 'tsan'],
                               platforms=['linux'],
                               labels=['sanitizers'],
-                              extra_args=extra_args)
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
   return test_jobs
 
-  
-def _create_portability_test_jobs(extra_args=[]):
+
+def _create_portability_test_jobs(extra_args=[], inner_jobs=_DEFAULT_INNER_JOBS):
   test_jobs = []
   # portability C x86
   test_jobs += _generate_jobs(languages=['c'],
@@ -159,8 +166,9 @@ def _create_portability_test_jobs(extra_args=[]):
                               arch='x86',
                               compiler='default',
                               labels=['portability'],
-                              extra_args=extra_args)
-  
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
+
   # portability C and C++ on x64
   for compiler in ['gcc4.4', 'gcc4.6', 'gcc5.3',
                    'clang3.5', 'clang3.6', 'clang3.7']:
@@ -170,8 +178,9 @@ def _create_portability_test_jobs(extra_args=[]):
                                 arch='x64',
                                 compiler=compiler,
                                 labels=['portability'],
-                                extra_args=extra_args)
-  
+                                extra_args=extra_args,
+                                inner_jobs=inner_jobs)
+
   # portability C on Windows
   for arch in ['x86', 'x64']:
     for compiler in ['vs2013', 'vs2015']:
@@ -181,24 +190,27 @@ def _create_portability_test_jobs(extra_args=[]):
                                   arch=arch,
                                   compiler=compiler,
                                   labels=['portability'],
-                                  extra_args=extra_args)
-  
+                                  extra_args=extra_args,
+                                  inner_jobs=inner_jobs)
+
   test_jobs += _generate_jobs(languages=['python'],
                               configs=['dbg'],
                               platforms=['linux'],
                               arch='default',
                               compiler='python3.4',
                               labels=['portability'],
-                              extra_args=extra_args)
-  
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
+
   test_jobs += _generate_jobs(languages=['csharp'],
                               configs=['dbg'],
                               platforms=['linux'],
                               arch='default',
                               compiler='coreclr',
                               labels=['portability'],
-                              extra_args=extra_args)
-  return test_jobs  
+                              extra_args=extra_args,
+                              inner_jobs=inner_jobs)
+  return test_jobs
 
 
 def _allowed_labels():
@@ -212,7 +224,7 @@ def _allowed_labels():
 
 argp = argparse.ArgumentParser(description='Run a matrix of run_tests.py tests.')
 argp.add_argument('-j', '--jobs',
-                  default=multiprocessing.cpu_count()/_INNER_JOBS,
+                  default=multiprocessing.cpu_count()/_DEFAULT_INNER_JOBS,
                   type=int,
                   help='Number of concurrent run_tests.py instances.')
 argp.add_argument('-f', '--filter',
@@ -241,15 +253,21 @@ argp.add_argument('--base_branch',
                   default='origin/master',
                   type=str,
                   help='Branch that pull request is requesting to merge into')
+argp.add_argument('--inner_jobs',
+                  default=_DEFAULT_INNER_JOBS,
+                  type=int,
+                  help='Number of jobs in each run_tests.py instance')
 args = argp.parse_args()
 
+
 extra_args = []
 if args.build_only:
   extra_args.append('--build_only')
 if args.force_default_poller:
   extra_args.append('--force_default_poller')
 
-all_jobs = _create_test_jobs(extra_args=extra_args) + _create_portability_test_jobs(extra_args=extra_args)
+all_jobs = _create_test_jobs(extra_args=extra_args, inner_jobs=args.inner_jobs) + \
+           _create_portability_test_jobs(extra_args=extra_args, inner_jobs=args.inner_jobs)
 
 jobs = []
 for job in all_jobs: