Bug 1309060 - Give |mach python-test| the ability to run tests in parallel, r=gps
authorAndrew Halberstadt <ahalberstadt@mozilla.com>
Tue, 11 Oct 2016 12:29:09 -0400
changeset 322685 083a65f9547140f0da619a17ef87f28bf8e36d8b
parent 322684 a0d07b344b2ffaa62218fb3c79fdcc69019240e7
child 322686 55a41c5e8e88010ea8a539c009f27b7b14988700
push id30960
push userkwierso@gmail.com
push dateThu, 17 Nov 2016 00:42:57 +0000
treeherdermozilla-central@830ce59e0a13 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersgps
bugs1309060
milestone53.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1309060 - Give |mach python-test| the ability to run tests in parallel, r=gps We recently switched make check to call into |mach python-test| rather than invoking python itself for each test file. But this ended up slowing down the tests as they were no longer being run in parallel. This patch adds a --jobs flag to python-tests and runs test files in parallel. Note: if more than one job is used, output per test will be buffered and printed at the end to avoid interleaving. This has the unfortunate side effect of making |mach python-test| look like it is hanging, especially if running a very long file like mozbase's test.py. For this reason, we still use -j1 by default so output will continue to be streamed. In automation we will use multiple processes though. MozReview-Commit-ID: 3u0wOFmyQLI
python/mach_commands.py
--- a/python/mach_commands.py
+++ b/python/mach_commands.py
@@ -1,20 +1,25 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
-import __main__
 import argparse
 import logging
 import mozpack.path as mozpath
 import os
 
+from concurrent.futures import (
+    ThreadPoolExecutor,
+    as_completed,
+    thread,
+)
+
 from mozbuild.base import (
     MachCommandBase,
 )
 
 from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
@@ -48,27 +53,32 @@ class MachCommands(MachCommandBase):
         default=False,
         action='store_true',
         help='Stop running tests after the first error or failure.')
     @CommandArgument('--path-only',
         default=False,
         action='store_true',
         help=('Collect all tests under given path instead of default '
               'test resolution. Supports pytest-style tests.'))
+    @CommandArgument('-j', '--jobs',
+        default=1,
+        type=int,
+        help='Number of concurrent jobs to run. Default is 1.')
     @CommandArgument('tests', nargs='*',
         metavar='TEST',
         help=('Tests to run. Each test can be a single file or a directory. '
               'Default test resolution relies on PYTHON_UNIT_TESTS.'))
     def python_test(self,
                     tests=[],
                     test_objects=None,
                     subsuite=None,
                     verbose=False,
                     path_only=False,
-                    stop=False):
+                    stop=False,
+                    jobs=1):
         self._activate_virtualenv()
 
         def find_tests_by_path():
             import glob
             files = []
             for t in tests:
                 if t.endswith('.py') and os.path.isfile(t):
                     files.append(t)
@@ -85,18 +95,16 @@ class MachCommands(MachCommandBase):
             return files
 
         # Python's unittest, and in particular discover, has problems with
         # clashing namespaces when importing multiple test modules. What follows
         # is a simple way to keep environments separate, at the price of
         # launching Python multiple times. Most tests are run via mozunit,
         # which produces output in the format Mozilla infrastructure expects.
         # Some tests are run via pytest.
-        return_code = 0
-        found_tests = False
         if test_objects is None:
             # If we're not being called from `mach test`, do our own
             # test resolution.
             if path_only:
                 if tests:
                     test_objects = [{'path': p} for p in find_tests_by_path()]
                 else:
                     self.log(logging.WARN, 'python-test', {},
@@ -108,51 +116,85 @@ class MachCommands(MachCommandBase):
                 if tests:
                     # If we were given test paths, try to find tests matching them.
                     test_objects = resolver.resolve_tests(paths=tests,
                                                           flavor='python')
                 else:
                     # Otherwise just run everything in PYTHON_UNIT_TESTS
                     test_objects = resolver.resolve_tests(flavor='python')
 
-        for test in test_objects:
-            found_tests = True
-            f = test['path']
-            file_displayed_test = []  # Used as a boolean.
-
-            def _line_handler(line):
-                if not file_displayed_test:
-                    output = ('Ran' in line or 'collected' in line or
-                              line.startswith('TEST-'))
-                    if output:
-                        file_displayed_test.append(True)
-
-            inner_return_code = self.run_process(
-                [self.virtualenv_manager.python_path, f],
-                ensure_exit_code=False,  # Don't throw on non-zero exit code.
-                log_name='python-test',
-                # subprocess requires native strings in os.environ on Windows
-                append_env={b'PYTHONDONTWRITEBYTECODE': str('1')},
-                line_handler=_line_handler)
-            return_code += inner_return_code
-
-            if not file_displayed_test:
-                self.log(logging.WARN, 'python-test', {'file': f},
-                         'TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() call?): {file}')
-
-            if verbose:
-                if inner_return_code != 0:
-                    self.log(logging.INFO, 'python-test', {'file': f},
-                             'Test failed: {file}')
-                else:
-                    self.log(logging.INFO, 'python-test', {'file': f},
-                             'Test passed: {file}')
-            if stop and return_code > 0:
-                return 1
-
-        if not found_tests:
+        if not test_objects:
             message = 'TEST-UNEXPECTED-FAIL | No tests collected'
             if not path_only:
-                 message += ' (Not in PYTHON_UNIT_TESTS? Try --path-only?)'
+                message += ' (Not in PYTHON_UNIT_TESTS? Try --path-only?)'
             self.log(logging.WARN, 'python-test', {}, message)
             return 1
 
-        return 0 if return_code == 0 else 1
+        self.jobs = jobs
+        self.terminate = False
+        self.verbose = verbose
+
+        return_code = 0
+        with ThreadPoolExecutor(max_workers=self.jobs) as executor:
+            futures = [executor.submit(self._run_python_test, test['path'])
+                       for test in test_objects]
+
+            try:
+                for future in as_completed(futures):
+                    output, ret = future.result()
+
+                    for line in output:
+                        self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
+
+                    return_code = return_code or ret
+            except KeyboardInterrupt:
+                # Hack to force stop currently running threads.
+                # https://gist.github.com/clchiou/f2608cbe54403edb0b13
+                executor._threads.clear()
+                thread._threads_queues.clear()
+                raise
+
+        return return_code
+
+    def _run_python_test(self, test_path):
+        from mozprocess import ProcessHandler
+
+        output = []
+
+        def _log(line):
+            # Buffer messages if more than one worker to avoid interleaving
+            if self.jobs > 1:
+                output.append(line)
+            else:
+                self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
+
+        file_displayed_test = []  # used as boolean
+
+        def _line_handler(line):
+            if not file_displayed_test:
+                output = ('Ran' in line or 'collected' in line or
+                          line.startswith('TEST-'))
+                if output:
+                    file_displayed_test.append(True)
+
+            _log(line)
+
+        _log(test_path)
+        cmd = [self.virtualenv_manager.python_path, test_path]
+        env = os.environ.copy()
+        env[b'PYTHONDONTWRITEBYTECODE'] = b'1'
+
+        proc = ProcessHandler(cmd, env=env, processOutputLine=_line_handler, storeOutput=False)
+        proc.run()
+
+        return_code = proc.wait()
+
+        if not file_displayed_test:
+            _log('TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() '
+                 'call?): {}'.format(test_path))
+
+        if self.verbose:
+            if return_code != 0:
+                _log('Test failed: {}'.format(test_path))
+            else:
+                _log('Test passed: {}'.format(test_path))
+
+        return output, return_code