[llvm] 38edab1 - [lit] Improve handling of parallelism group semaphores

[Julian Lettner via llvm-commits]

Author: Julian Lettner
Date: 2020-04-06T20:52:06-07:00
New Revision: 38edab1c400e74783930cb542ff870ee8cd5de91

URL: https://github.com/llvm/llvm-project/commit/38edab1c400e74783930cb542ff870ee8cd5de91
DIFF: https://github.com/llvm/llvm-project/commit/38edab1c400e74783930cb542ff870ee8cd5de91.diff

LOG: [lit] Improve handling of parallelism group semaphores

Added: 
    

Modified: 
    llvm/utils/lit/lit/run.py
    llvm/utils/lit/lit/worker.py
    llvm/utils/lit/tests/Inputs/parallelism-groups/lit.cfg
    llvm/utils/lit/tests/parallelism-groups.py

Removed: 
    


################################################################################
diff  --git a/llvm/utils/lit/lit/run.py b/llvm/utils/lit/lit/run.py
index fca3fb47fd5e..212e909dcea1 100644
--- a/llvm/utils/lit/lit/run.py
+++ b/llvm/utils/lit/lit/run.py
@@ -7,13 +7,6 @@
 import lit.worker
 
 
-# No-operation semaphore for supporting `None` for parallelism_groups.
-#   lit_config.parallelism_groups['my_group'] = None
-class NopSemaphore(object):
-    def acquire(self): pass
-    def release(self): pass
-
-
 class MaxFailuresError(Exception):
     pass
 class TimeoutError(Exception):
@@ -66,13 +59,12 @@ def execute(self):
                 test.setResult(lit.Test.Result(lit.Test.UNRESOLVED, '', 0.0))
 
     def _execute(self, deadline):
-        semaphores = {
-            k: NopSemaphore() if v is None else
-            multiprocessing.BoundedSemaphore(v) for k, v in
-            self.lit_config.parallelism_groups.items()}
-
         self._increase_process_limit()
 
+        semaphores = {k: multiprocessing.BoundedSemaphore(v)
+                      for k, v in self.lit_config.parallelism_groups.items()
+                      if v is not None}
+
         pool = multiprocessing.Pool(self.workers, lit.worker.initialize,
                                     (self.lit_config, semaphores))
 

diff  --git a/llvm/utils/lit/lit/worker.py b/llvm/utils/lit/lit/worker.py
index 95a8578afac9..04fc77239e02 100644
--- a/llvm/utils/lit/lit/worker.py
+++ b/llvm/utils/lit/lit/worker.py
@@ -5,6 +5,7 @@
 For efficiency, we copy all data needed to execute all tests into each worker
 and store it in global variables. This reduces the cost of each task.
 """
+import contextlib
 import signal
 import time
 import traceback
@@ -39,28 +40,27 @@ def execute(test):
     Arguments and results of this function are pickled, so they should be cheap
     to copy.
     """
-    result = _execute_in_parallelism_group(test, _lit_config,
-                                           _parallelism_semaphores)
+    with _get_parallelism_semaphore(test):
+        result = _execute(test, _lit_config)
+
     test.setResult(result)
     return test
 
 
-def _execute_in_parallelism_group(test, lit_config, parallelism_semaphores):
+# TODO(python3): replace with contextlib.nullcontext
+ at contextlib.contextmanager
+def NopSemaphore():
+    yield
+
+
+def _get_parallelism_semaphore(test):
     pg = test.config.parallelism_group
     if callable(pg):
         pg = pg(test)
-
-    if pg:
-        semaphore = parallelism_semaphores[pg]
-        try:
-            semaphore.acquire()
-            return _execute(test, lit_config)
-        finally:
-            semaphore.release()
-    else:
-        return _execute(test, lit_config)
+    return _parallelism_semaphores.get(pg, NopSemaphore())
 
 
+# Do not inline! Directly used by LitTestCase.py
 def _execute(test, lit_config):
     start = time.time()
     result = _execute_test_handle_errors(test, lit_config)

diff  --git a/llvm/utils/lit/tests/Inputs/parallelism-groups/lit.cfg b/llvm/utils/lit/tests/Inputs/parallelism-groups/lit.cfg
index ffc41b1ffe54..de511877ee00 100644
--- a/llvm/utils/lit/tests/Inputs/parallelism-groups/lit.cfg
+++ b/llvm/utils/lit/tests/Inputs/parallelism-groups/lit.cfg
@@ -9,3 +9,7 @@ config.test_exec_root = None
 lit_config.parallelism_groups['my_group'] = None
 
 config.parallelism_group = 'my_group'
+
+# TODO(yln): we should have a nice property setter that doesn't allow setting
+# to a non-existent group and do Math.min(old_group, new_group) when we
+# overwrite an existing group.

diff  --git a/llvm/utils/lit/tests/parallelism-groups.py b/llvm/utils/lit/tests/parallelism-groups.py
index d80f2318fe05..bd0cf37ab0f0 100644
--- a/llvm/utils/lit/tests/parallelism-groups.py
+++ b/llvm/utils/lit/tests/parallelism-groups.py
@@ -9,9 +9,6 @@
 # [project/lit.cfg]
 #   config.parallelism_group = 'my_group'
 #
-# Note: We need at least 2 tests to prevent lit from using "single process
-#       mode", which ignores parallelism groups.
-#
 
 # RUN: %{lit} -j2 %{inputs}/parallelism-groups | FileCheck %s