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

[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.
----------------
prasoon054 wrote:

Yes. I'll add links and a TODO. The cancel hazard is [cpython#94440](https://github.com/python/cpython/issues/94440). Cancelling a future can cause shutdown to hang. On Python 3.8, I also observed the management thread crashing with `InvalidStateError` when it attempted to set `BrokenProcessPool` on futures that had already been cancelled. This was fixed upstream and backported to Python 3.10 and 3.11, but not to 3.8, so the workarouond is required as long as our minimum supported version remains below 3.10.

The long-term replacement would be [`ProcessPoolExecutor.kill_workers()`](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor.kill_workers), introduced in Python 3.14

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