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

[Prasoon Kumar via llvm-commits]

================
@@ -103,59 +166,111 @@ 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 = {}
 
         try:
-            self._wait_for(async_results, deadline)
-        except:
-            # Terminate all pools on exception
-            for pool in pools:
-                pool.terminate()
+            self._dispatch_and_wait(executors, future_to_test, deadline)
+        except BaseException:
+            self._abort_executors(executors, future_to_test)
             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:
-            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():
-                    self.failures += 1
-                    if self.failures == self.max_failures:
-                        raise MaxFailuresError()
-            idx += 1
+        else:
+            for ex in executors:
+                # On macOS, Queue.join_thread() inside shutdown(wait=True)
+                # deadlocks: join_executor_internals() calls it before
+                # p.join(), but macOS requires the inverse order.
+                # cancel_join_thread() makes join_thread() a no-op;
+                # the feeder still delivers sentinels before the write end
+                # closes.
+                if hasattr(ex, "_call_queue") and ex._call_queue is not None:
+                    ex._call_queue.cancel_join_thread()
+                ex.shutdown(wait=True)
+
+    def _dispatch_and_wait(self, executors, future_to_test, deadline):
+        """Submits tests to executors and collects results as they complete.
+
+        Bounds the number of futures outstanding at any time to at most
+        window (see SUBMISSION_WINDOW_PER_WORKER), submitting exactly one
+        new test for each one that completes. Submitting every test up
+        front floods the executor's wakeup pipe and can deadlock submit()
+        against the executor's manager thread on Python <= 3.11.5
+        (https://github.com/python/cpython/issues/105829)
+
+        Mutates future_to_test in place: adds an entry for every test
+        submitted, and removes it once that test's result has been
+        collected. On return, or if this call raises, future_to_test
+        holds exactly the futures that have not yet been collected, which
+        the caller's abort path relies on.
+
+        Args:
+            executors: The ProcessPoolExecutor pool(s) tests are dispatched to.
+            future_to_test: A dict mapping each in-flight Future to its
+              corresponding Test. Populated and drained by this call.
+            deadline: The absolute time (as returned by time.time()) after
+              which the call raises TimeoutError.
+
+        Raises:
+            TimeoutError: deadline passed with the tests still outstanding.
+            MaxFailuresError: The number of failed tests reached self.max_failures.
+            WorkerCrashError: A worker process died unexpectedly (e.g.
+              segfault, OOM-kill) instead of returning a result.
+        """
+        try:
+            window = int(
+                os.getenv(
+                    "LIT_SUBMISSION_WINDOW",
+                    SUBMISSION_WINDOW_PER_WORKER * self.workers,
+                )
+            ) or len(self.tests)
+            tests_iter = enumerate(self.tests)
+            pending = set()
+
+            def submit_next():
+                """Submits the next not-yet-submitted test, if any.
+
+                Returns:
+                    True if a test was submitted, False if none remained.
+                """
+                for i, test in tests_iter:
+                    ex = executors[i % len(executors)]
+                    future = ex.submit(lit.worker.execute, test)
+                    future_to_test[future] = test
+                    pending.add(future)
+                    return True
+                return False
+
+            while len(pending) < window and submit_next():
+                pass
+
+            while pending:
+                done, pending = wait(
+                    pending,
+                    timeout=deadline - time.time(),
+                    return_when=FIRST_COMPLETED,
+                )
+                if not done:
+                    raise TimeoutError()
+                for future in done:
+                    remote_test = future.result()
+                    local_test = future_to_test.pop(future)
+                    self._update_test(local_test, remote_test)
+                    self.progress_callback(remote_test)
+                    if remote_test.isFailure():
+                        self.failures += 1
+                        if self.failures == self.max_failures:
----------------
prasoon054 wrote:

`==` is safe here since it expects a positive integer. https://github.com/llvm/llvm-project/blob/4a75259134c72d8be7fa3ea36c9dc7bfc02640ac/llvm/utils/lit/lit/cl_arguments.py#L345-L353 
But `>=` reads better on its own since it doesn't depend on knowing `max_failure` can't be 0, so I'll switch to that.

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