simulate.py: Run tests in parallel

util/sim/simulate.py

@@ -8,11 +8,14 @@
 # TODO colluca: timeout feature
 
 import argparse
+import multiprocessing
 from pathlib import Path
 import subprocess
 from termcolor import colored, cprint
+import os
 import re
 import sys
+import time
 import yaml
 
 
@@ -62,6 +65,16 @@ def parse_args():
         '--early-exit',
         action='store_true',
         help='Exit as soon as any test fails')
+    parser.add_argument(
+        '-j',
+        action='store',
+        dest='n_procs',
+        nargs='?',
+        default=1,
+        const=os.cpu_count(),
+        help=('Maximum number of tests to run in parallel. '
+              'One if the option is not present. Equal to the number of CPU cores '
+              'if the option is present but not followed by an argument.'))
     args = parser.parse_args()
     return args
 
@@ -81,17 +94,19 @@ def check_exit_code(test, exit_code):
         return exit_code
 
 
-def run_simulation(cmd, simulator, test):
+def run_simulation(cmd, simulator, test, quiet=False):
     # Defaults
     result = 1
 
     # Spawn simulation subprocess
-    p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, universal_newlines=True)
+    p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                         universal_newlines=True)
 
     # Poll simulation subprocess and log its output
     while p.poll() is None:
         line = p.stdout.readline()
-        print(line, end='', flush=True)
+        if not quiet:
+            print(line, end='', flush=True)
 
         # When simulating with vsim or vcs, we need to parse the simulation
         # log to catch the application's return code
@@ -132,6 +147,7 @@ def run_test(test, args):
     sim_bin = args.sim_bin if args.sim_bin else SIMULATOR_BINS[simulator]
     dry_run = args.dry_run
     testlist = args.testlist
+    quiet = args.n_procs != 1
 
     # Check if simulator is supported for this test
     if 'simulators' in test:
@@ -152,12 +168,13 @@ def run_test(test, args):
     else:
         cmd = SIMULATOR_CMDS[simulator]
         cmd = cmd.format(sim_bin=sim_bin, elf=elf)
-    print(f'$ {cmd}', flush=True)
+    if not quiet:
+        print(f'$ {cmd}', flush=True)
 
     # Run simulation
     result = 0
     if not dry_run:
-        result = run_simulation(cmd, simulator, test)
+        result = run_simulation(cmd, simulator, test, quiet)
 
     # Report failure or success
     if result != 0:
@@ -172,32 +189,57 @@ def print_failed_test(test):
     print(f'{colored(test["elf"], "cyan")} test {colored("failed", "red")}')
 
 
-def print_test_summary(failed_tests, dry_run=False):
-    if not dry_run:
-        print('\n==== Test summary ====')
+def print_test_summary(failed_tests, args):
+    if not args.dry_run:
+        header = f'\n==== Test summary {"(early exit)" if args.early_exit else ""} ===='
+        cprint(header, attrs=['bold'])
         if failed_tests:
             for failed_test in failed_tests:
                 print_failed_test(failed_test)
-            return 1
         else:
             print(f'{colored("All tests passed!", "green")}')
-            return 0
-    return 0
 
 
 def run_tests(args):
-    # Iterate tests
+
+    # Get tests from testlist
     tests = get_tests(args.testlist)
+
+    # Create a process Pool
+    with multiprocessing.Pool(args.n_procs) as pool:
+
+        # Create a shared object which parent and child processes can access
+        # concurrently to terminate the pool early as soon as one process fails
+        exit_early = multiprocessing.Value('B')
+        exit_early.value = 0
+
+        # Define callback for early exit
+        def completion_callback(result):
+            if args.early_exit and result != 0:
+                exit_early.value = 1
+
+        # Queue tests to process pool
+        results = []
+        for test in tests:
+            result = pool.apply_async(run_test, args=(test, args), callback=completion_callback)
+            results.append(result)
+
+        # Wait for all tests to complete
+        running = range(len(tests))
+        while len(running) != 0 and not exit_early.value:
+            time.sleep(1)
+            running = [i for i in running if not results[i].ready()]
+
+
+    # Query test results
     failed_tests = []
-    for test in tests:
-        # Run test
-        result = run_test(test, args)
-        if result != 0:
+    for test, result in zip(tests, results):
+        if result.ready() and result.get() != 0:
             failed_tests.append(test)
-            # End program if requested on first test failure
-            if args.early_exit:
-                break
-    return print_test_summary(failed_tests, args.dry_run)
+
+    print_test_summary(failed_tests, args)
+
+    return len(failed_tests)
 
 
 def main():