[lit][NFC] Cleanup lit worker process handling

[yln]

Summary

Move code that is executed on worker process to separate file. This
makes the use of the pickled arguments stored in global variables in the
worker a bit clearer. (Still not pretty though.)

Extract handling of parallelism groups to it's own function.

Use BoundedSemaphore instead of Semaphore. BoundedSemaphore raises for
unmatched release() calls.

Cleanup imports.

[rnk]

llvm/utils/lit/lit/run.py

166

Won't we get an unbound local error now if parallelism_semaphores[pg] raises KeyError? I would expect this to raise, hit the except block, then run the finally, where it would raise again.

201

I agree this check seems unnecessary. It seems like it would be a bug to insert None as a value in the parallelism_semaphores dict.

[yln]

llvm/utils/lit/lit/run.py

166

Yes, it raises, prints a stack trace, and crashes. (Instead of marking tests as unresolved when I misconfigure my parallelism groups.)

201

Your comment about None gave me an idea. I think it *would* be useful for intentionally stating "no restrictions" vs "unsure if configuration error" or just not concerned about parallelism. I will update the patch. Thanks!

[yln]

Widened scope of cleanup.

[yln]

@rnk
I just realized that we are still using "threads" for the display and command line options and different notions in the code: thread, jobs, workers. Since I am already at it, should I unify the terminology (in a separate patch)?

[rnk]

In D58196#1398802, @yln wrote:

@rnk
I just realized that we are still using "threads" for the display and command line options and different notions in the code: thread, jobs, workers. Since I am already at it, should I unify the terminology (in a separate patch)?

Sure. The way I see it "worker" is a more general term, so it's reasonable to use it in comments even if we use the --numThreads command line option.

llvm/utils/lit/lit/worker.py

51–53

This code is 2-space indented, but the rest of lit is mostly 4-space.

[yln]

Fix indentation.

[rnk]

lgtm

llvm/utils/lit/lit/run.py

-import os
-import sys
-import threading
+import multiprocessing
 import time
-import traceback
-try:
-    import Queue as queue
-except ImportError:
-    import queue
-
-try:
-    import win32api
-except ImportError:
-    win32api = None
 
-import multiprocessing
 import lit.Test
-
-def abort_now():
-    """Abort the current process without doing any exception teardown"""
-    sys.stdout.flush()
-    if win32api:
-        win32api.TerminateProcess(win32api.GetCurrentProcess(), 3)
-    else:
-        os.kill(0, 9)
+import lit.util
+import lit.worker
 
 class _Display(object):
     def __init__(self, display, provider, maxFailures):
         self.display = display
         self.provider = provider
         self.maxFailures = maxFailures or object()
         self.failedCount = 0
     def update(self, test):
class Run(object):
 
     def __init__(self, lit_config, tests):
         self.lit_config = lit_config
         self.tests = tests
         # Set up semaphores to limit parallelism of certain classes of tests.
         # For example, some ASan tests require lots of virtual memory and run
         # faster with less parallelism on OS X.
         self.parallelism_semaphores = \
-                {k: multiprocessing.Semaphore(v) for k, v in
+                {k: multiprocessing.BoundedSemaphore(v) for k, v in
                  self.lit_config.parallelism_groups.items()}
 
     def execute_test(self, test):
-        return _execute_test_impl(test, self.lit_config,
-                                  self.parallelism_semaphores)
+        return lit.worker._execute_test(test, self.lit_config)
 
     def execute_tests_in_pool(self, jobs, max_time):
         # We need to issue many wait calls, so compute the final deadline and
         # subtract time.time() from that as we go along.
         deadline = None
         if max_time:
             deadline = time.time() + max_time
 
         # Start a process pool. Copy over the data shared between all test runs.
         # FIXME: Find a way to capture the worker process stderr. If the user
         # interrupts the workers before we make it into our task callback, they
         # will each raise a KeyboardInterrupt exception and print to stderr at
         # the same time.
-        pool = multiprocessing.Pool(jobs, worker_initializer,
+        pool = multiprocessing.Pool(jobs, lit.worker.initializer,
                                     (self.lit_config,
                                      self.parallelism_semaphores))
 
         # Install a console-control signal handler on Windows.
-        if win32api is not None:
+        if lit.util.win32api is not None:
             def console_ctrl_handler(type):
                 print('\nCtrl-C detected, terminating.')
                 pool.terminate()
                 pool.join()
-                abort_now()
+                lit.util.abort_now()
                 return True
-            win32api.SetConsoleCtrlHandler(console_ctrl_handler, True)
+            lit.util.win32api.SetConsoleCtrlHandler(console_ctrl_handler, True)
 
         try:
-            async_results = [pool.apply_async(worker_run_one_test,
+            async_results = [pool.apply_async(lit.worker.run_one_test,
                                               args=(test_index, test),
                                               callback=self.consume_test_result)
                              for test_index, test in enumerate(self.tests)]
             pool.close()
 
             # Wait for all results to come in. The callback that runs in the
             # parent process will update the display.
             for a in async_results:
def execute_tests(self, display, jobs, max_time=None):
 
         # Save the display object on the runner so that we can update it from
         # our task completion callback.
         self.display = display
 
         self.failure_count = 0
         self.hit_max_failures = False
         if jobs == 1:
-            global child_lit_config
-            child_lit_config = self.lit_config
             for test_index, test in enumerate(self.tests):
-                result = worker_run_one_test(test_index, test)
-                self.consume_test_result(result)
+                lit.worker._execute_test(test, self.lit_config)
+                self.consume_test_result((test_index, test))
                 if self.hit_max_failures:
                     break
         else:
             self.execute_tests_in_pool(jobs, max_time)
 
         # Mark any tests that weren't run as UNRESOLVED.
         for test in self.tests:
             if test.result is None:
                 test.setResult(lit.Test.Result(lit.Test.UNRESOLVED, '', 0.0))
 
     def consume_test_result(self, pool_result):
-        """Test completion callback for worker_run_one_test
+        """Test completion callback for lit.worker.run_one_test
 
         Updates the test result status in the parent process. Each task in the
         pool returns the test index and the result, and we use the index to look
         up the original test object. Also updates the progress bar as tasks
         complete.
         """
         # Don't add any more test results after we've hit the maximum failure
         # count.  Otherwise we're racing with the main thread, which is going
def consume_test_result(self, pool_result):
         self.tests[test_index] = test_with_result
         self.display.update(test_with_result)
 
         # If we've finished all the tests or too many tests have failed, notify
         # the main thread that we've stopped testing.
         self.failure_count += (test_with_result.result.code == lit.Test.FAIL)
         if self.lit_config.maxFailures and \
                 self.failure_count == self.lit_config.maxFailures:
             self.hit_max_failures = True
-
-def _execute_test_impl(test, lit_config, parallelism_semaphores):
-    """Execute one test"""
-    pg = test.config.parallelism_group
-    if callable(pg):
-        pg = pg(test)
-
-    result = None
-    semaphore = None
-    try:
-        if pg:
-            semaphore = parallelism_semaphores[pg]
-        if semaphore:
-            semaphore.acquire()
-        start_time = time.time()
-        result = test.config.test_format.execute(test, lit_config)
-        # Support deprecated result from execute() which returned the result
-        # code and additional output as a tuple.
-        if isinstance(result, tuple):
-            code, output = result
-            result = lit.Test.Result(code, output)
-        elif not isinstance(result, lit.Test.Result):
-            raise ValueError("unexpected result from test execution")
-        result.elapsed = time.time() - start_time
-    except KeyboardInterrupt:
-        raise
-    except:
-        if lit_config.debug:
-            raise
-        output = 'Exception during script execution:\n'
-        output += traceback.format_exc()
-        output += '\n'
-        result = lit.Test.Result(lit.Test.UNRESOLVED, output)
-    finally:
-        if semaphore:
-            semaphore.release()
-
-    test.setResult(result)
-
-child_lit_config = None
-child_parallelism_semaphores = None
-
-def worker_initializer(lit_config, parallelism_semaphores):
-    """Copy expensive repeated data into worker processes"""
-    global child_lit_config
-    child_lit_config = lit_config
-    global child_parallelism_semaphores
-    child_parallelism_semaphores = parallelism_semaphores
-
-def worker_run_one_test(test_index, test):
-    """Run one test in a multiprocessing.Pool
-
-    Side effects in this function and functions it calls are not visible in the
-    main lit process.
-
-    Arguments and results of this function are pickled, so they should be cheap
-    to copy. For efficiency, we copy all data needed to execute all tests into
-    each worker and store it in the child_* global variables. This reduces the
-    cost of each task.
-
-    Returns an index and a Result, which the parent process uses to update
-    the display.
-    """
-    try:
-        _execute_test_impl(test, child_lit_config, child_parallelism_semaphores)
-        return (test_index, test)
-    except KeyboardInterrupt as e:
-        # If a worker process gets an interrupt, abort it immediately.
-        abort_now()
-    except:
-        traceback.print_exc()

llvm/utils/lit/lit/util.py

try:
         for child in children_iterator:
             try:
                 child.kill()
             except psutil.NoSuchProcess:
                 pass
         psutilProc.kill()
     except psutil.NoSuchProcess:
         pass
+
+
+try:
+    import win32api
+except ImportError:
+    win32api = None
+
+def abort_now():
+    """Abort the current process without doing any exception teardown"""
+    sys.stdout.flush()
+    if win32api:
+        win32api.TerminateProcess(win32api.GetCurrentProcess(), 3)
+    else:
+        os.kill(0, 9)

llvm/utils/lit/lit/worker.py

+# The functions in this module are meant to run on a separate worker process.
+# Exception: in single process mode _execute_test is called directly.
+import time
+import traceback
+
+import lit.Test
+import lit.util
+
+_lit_config = None
+_parallelism_semaphores = None
+
+def initializer(lit_config, parallelism_semaphores):
+    """Copy expensive repeated data into worker processes"""
+    global _lit_config
+    global _parallelism_semaphores
+    _lit_config = lit_config
+    _parallelism_semaphores = parallelism_semaphores
+
+def run_one_test(test_index, test):
+    """Run one test in a multiprocessing.Pool
+
+    Side effects in this function and functions it calls are not visible in the
+    main lit process.
+
+    Arguments and results of this function are pickled, so they should be cheap
+    to copy. For efficiency, we copy all data needed to execute all tests into
+    each worker and store it in the worker_* global variables. This reduces the
+    cost of each task.
+
+    Returns an index and a Result, which the parent process uses to update
+    the display.
+    """
+    try:
+        _execute_test_in_parallelism_group(test, _lit_config,
+                                           _parallelism_semaphores)
+        return (test_index, test)
+    except KeyboardInterrupt:
+        # If a worker process gets an interrupt, abort it immediately.
+        lit.util.abort_now()
+    except:
+        traceback.print_exc()
+
+def _execute_test_in_parallelism_group(test, lit_config, parallelism_semaphores):
+    """Execute one test inside the appropriate parallelism group"""
+    pg = test.config.parallelism_group
+    if callable(pg):
+        pg = pg(test)
+
+    if pg:
+        semaphore = parallelism_semaphores[pg]
+        try:
+            semaphore.acquire()
+            _execute_test(test, lit_config)
+        finally:
+            semaphore.release()
+    else:
+        _execute_test(test, lit_config)
+
+def _execute_test(test, lit_config):
+    """Execute one test"""
+    try:
+        start_time = time.time()
+        result = test.config.test_format.execute(test, lit_config)
+        # Support deprecated result from execute() which returned the result
+        # code and additional output as a tuple.
+        if isinstance(result, tuple):
+            code, output = result
+            result = lit.Test.Result(code, output)
+        elif not isinstance(result, lit.Test.Result):
+            raise ValueError("unexpected result from test execution")
+        result.elapsed = time.time() - start_time
+    except KeyboardInterrupt:
+        raise
+    except:
+        if lit_config.debug:
+            raise
+        output = 'Exception during script execution:\n'
+        output += traceback.format_exc()
+        output += '\n'
+        result = lit.Test.Result(lit.Test.UNRESOLVED, output)
+
+    test.setResult(result)