[llvm] [lit] Migrate lit to ProcessPoolExecutor (PR #202681)

[Aiden Grossman via llvm-commits]

================
@@ -103,59 +113,62 @@ def _execute(self, deadline):
                 % (num_pools, self.workers, workers_per_pool_list)
             )
 
-        # Create multiple pools
-        pools = []
-        for pool_size in workers_per_pool_list:
-            pool = multiprocessing.Pool(
-                pool_size, lit.worker.initialize, (self.lit_config, semaphores)
+        executors = [
+            ProcessPoolExecutor(
+                max_workers=pool_size,
+                initializer=lit.worker.initialize,
+                initargs=(self.lit_config, semaphores),
             )
-            pools.append(pool)
-
-        # Distribute tests across pools
-        tests_per_pool = _ceilDiv(len(self.tests), num_pools)
-        async_results = []
-
-        for pool_idx, pool in enumerate(pools):
-            start_idx = pool_idx * tests_per_pool
-            end_idx = min(start_idx + tests_per_pool, len(self.tests))
-            for test in self.tests[start_idx:end_idx]:
-                ar = pool.apply_async(
-                    lit.worker.execute, args=[test], callback=self.progress_callback
-                )
-                async_results.append(ar)
+            for pool_size in workers_per_pool_list
+        ]
 
-        # Close all pools
-        for pool in pools:
-            pool.close()
+        future_to_test = {}
+        for i, test in enumerate(self.tests):
+            ex = executors[i % len(executors)]
+            future_to_test[ex.submit(lit.worker.execute, test)] = test
 
         try:
-            self._wait_for(async_results, deadline)
-        except:
-            # Terminate all pools on exception
-            for pool in pools:
-                pool.terminate()
-            raise
-        finally:
-            # Join all pools
-            for pool in pools:
-                pool.join()
-
-    def _wait_for(self, async_results, deadline):
-        timeout = deadline - time.time()
-        idx = 0
-        while len(async_results) > 0:
+            self._wait_for(future_to_test, deadline)
+        except BaseException:
             try:
-                ar = async_results.pop(0)
-                test = ar.get(timeout)
-            except multiprocessing.TimeoutError:
-                raise TimeoutError()
-            else:
-                self._update_test(self.tests[idx], test)
-                if test.isFailure():
+                tree_kill_ok, _ = lit.util.killProcessAndChildrenIsSupported()
+
+                # Prevent multiprocessing from joining queue feeder threads at exit
+                for ex in executors:
+                    if hasattr(ex, "_call_queue") and ex._call_queue is not None:
+                        ex._call_queue.cancel_join_thread()
+
+                for ex in executors:
+                    for pid, proc in list((ex._processes or {}).items()):
+                        if tree_kill_ok:
+                            lit.util.killProcessAndChildren(pid)
+                        else:
+                            proc.kill()
+
+                # Skip future.cancel() and ex.shutdown() to prevent
+                # ProcessPoolExecutor deadlocks and InvalidStateErrors on Python 3.8.
----------------
boomanaiden154 wrote:

Is there a python bug link or something that should go here?

Maybe also a TODO to remove the workaround if it's been fixed in more recent versions?

https://github.com/llvm/llvm-project/pull/202681