Add InterpreterPoolExecutor.

python · ericsnowcurrently · Mar 1, 2024 · Oct 20, 2023 · Mar 1, 2024 · Mar 1, 2024
commit 9774ffd5ea47bbbd200e6454579b9dc87a94711f
diff --git a/Lib/concurrent/futures/__init__.py b/Lib/concurrent/futures/__init__.py
@@ -30,6 +30,7 @@
     'as_completed',
     'ProcessPoolExecutor',
     'ThreadPoolExecutor',
+    'InterpreterPoolExecutor',
 )
 
 
@@ -38,7 +39,7 @@ def __dir__():
 
 
 def __getattr__(name):
-    global ProcessPoolExecutor, ThreadPoolExecutor
+    global ProcessPoolExecutor, ThreadPoolExecutor, InterpreterPoolExecutor
 
     if name == 'ProcessPoolExecutor':
         from .process import ProcessPoolExecutor as pe
@@ -50,4 +51,9 @@ def __getattr__(name):
         ThreadPoolExecutor = te
         return te
 
+    if name == 'InterpreterPoolExecutor':
+        from .interpreter import InterpreterPoolExecutor as ie
+        InterpreterPoolExecutor = ie
+        return ie
+
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,50 @@
+
+"""Implements InterpreterPoolExecutor."""
+
+import concurrent.futures.thread as _thread
+import pickle
+from test.support import interpreters
+
+
+class InterpreterPoolExecutor(_thread.ThreadPoolExecutor):
+
+    @classmethod
+    def _normalize_initializer(cls, initializer, initargs):
+        if initializer is None:
+            return initializer, initargs
+
+        pickled = pickle.dumps((initializer, initargs))
+        def initializer(ctx):
+            interp, _ = ctx
+            interp.exec(f'''if True:
+                import pickle
+                initializer, initargs = pickle.loads({pickled!r})
+                initializer(*initargs)
+                ''')
+        return initializer, ()
+
+    @classmethod
+    def _normalize_task(cls, fn, args, kwargs):
+        pickled = pickle.dumps((fn, args, kwargs))
+        def wrapped(ctx):
+            interp, results = ctx
+            interp.exec(f'''if True:
+                import pickle
+                fn, args, kwargs = pickle.loads({pickled!r})
+                res = fn(*args, **kwargs)
+                _results.put(res)
+                ''')
+            return results.get_nowait()
+        return wrapped, (), {}
+
+    @classmethod
+    def _run_worker(cls, *args):
+        interp = interpreters.create()
+        interp.exec('import test.support.interpreters.queues')
+        results = interpreters.create_queue()
+        interp.prepare_main(_results=results)
+        ctx = (interp, results)
+        try:
+            _thread._worker(ctx, *args)
+        finally:
+            interp.close()
diff --git a/Lib/test/test_concurrent_futures/executor.py b/Lib/test/test_concurrent_futures/executor.py
@@ -26,6 +26,10 @@ def make_dummy_object(_):
 class ExecutorTest:
     # Executor.shutdown() and context manager usage is tested by
     # ExecutorShutdownTest.
+
+    def _assertRaises(self, exctype, *args, **kwargs):
+        return self.assertRaises(exctype, *args, **kwargs)
+
     def test_submit(self):
         future = self.executor.submit(pow, 2, 8)
         self.assertEqual(256, future.result())
@@ -53,7 +57,7 @@ def test_map_exception(self):
         i = self.executor.map(divmod, [1, 1, 1, 1], [2, 3, 0, 5])
         self.assertEqual(i.__next__(), (0, 1))
         self.assertEqual(i.__next__(), (0, 1))
-        self.assertRaises(ZeroDivisionError, i.__next__)
+        self._assertRaises(ZeroDivisionError, i.__next__)
 
     @support.requires_resource('walltime')
     def test_map_timeout(self):

diff --git a/Lib/test/test_concurrent_futures/test_interpreter_pool.py b/Lib/test/test_concurrent_futures/test_interpreter_pool.py
@@ -0,0 +1,104 @@
+import contextlib
+from test.support import interpreters
+import os
+import unittest
+
+from concurrent.futures import InterpreterPoolExecutor
+from .executor import ExecutorTest, mul
+from .util import BaseTestCase, InterpreterPoolMixin, setup_module
+
+
+def wait_for_token(queue):
+    while True:
+        try:
+            queue.get(timeout=0.1)
+        except interpreters.QueueEmpty:
+            continue
+        break
+
+
+def log_n_wait(args):
+    logqueue, waitqueue, ident = args
+    logqueue.put_nowait(f"{ident=} started")
+    try:
+        wait_for_token(waitqueue)
+    finally:
+        logqueue.put_nowait(f"{ident=} stopped")
+
+
+class InterpreterPoolExecutorTest(InterpreterPoolMixin, ExecutorTest, BaseTestCase):
+
+    def _assertRaises(self, exctype, *args, **kwargs):
+        return self.assertRaises(interpreters.ExecutionFailed, *args, **kwargs)
+
+    def test_default_workers(self):
+        executor = InterpreterPoolExecutor()
+        expected = min(32, (os.process_cpu_count() or 1) + 4)
+        self.assertEqual(executor._max_workers, expected)
+
+    def test_saturation(self):
+        executor = InterpreterPoolExecutor(max_workers=4)
+        waitqueue = interpreters.create_queue(syncobj=True)
+
+        for i in range(15 * executor._max_workers):
+            executor.submit(wait_for_token, waitqueue)
+        self.assertEqual(len(executor._threads), executor._max_workers)
+        for i in range(15 * executor._max_workers):
+            waitqueue.put(None)
+        executor.shutdown(wait=True)
+
+    def test_idle_thread_reuse(self):
+        executor = InterpreterPoolExecutor()
+        executor.submit(mul, 21, 2).result()
+        executor.submit(mul, 6, 7).result()
+        executor.submit(mul, 3, 14).result()
+        self.assertEqual(len(executor._threads), 1)
+        executor.shutdown(wait=True)
+
+    def test_executor_map_current_future_cancel(self):
+        logqueue = interpreters.create_queue(syncobj=True)
+        waitqueue = interpreters.create_queue(syncobj=True)
+        idents = ['first', 'second', 'third']
+        _idents = iter(idents)
+
+        with InterpreterPoolExecutor(max_workers=1) as pool:
+            # submit work to saturate the pool
+            fut = pool.submit(log_n_wait, (logqueue, waitqueue, next(_idents)))
+            try:
+                with contextlib.closing(
+                    pool.map(log_n_wait,
+                             [(logqueue, waitqueue, ident) for ident in _idents],
+                             timeout=0)
+                ) as gen:
+                    with self.assertRaises(TimeoutError):
+                        next(gen)
+            finally:
+                for i, ident in enumerate(idents, 1):
+                    waitqueue.put_nowait(None)
+                #for ident in idents:
+                #    waitqueue.put_nowait(None)
+            fut.result()
+
+            # When the pool shuts down (due to the context manager),
+            # each worker's interpreter will be finalized.  When that
+            # happens, each item in the queue owned by each finalizing
+            # interpreter will be removed.  Thus, we must copy the
+            # items *before* leaving the context manager.
+            # XXX This can be surprising.  Perhaps give users
+            # the option to keep objects beyond the original interpreter?
+            assert not logqueue.empty(), logqueue.qsize()
+            log = []
+            while not logqueue.empty():
+                log.append(logqueue.get_nowait())
+
+        # ident='second' is cancelled as a result of raising a TimeoutError
+        # ident='third' is cancelled because it remained in the collection of futures
+        self.assertListEqual(log, ["ident='first' started", "ident='first' stopped"])
+
+
+def setUpModule():
+    setup_module()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/Lib/test/test_concurrent_futures/util.py b/Lib/test/test_concurrent_futures/util.py
@@ -74,6 +74,10 @@ class ThreadPoolMixin(ExecutorMixin):
     executor_type = futures.ThreadPoolExecutor
 
 
+class InterpreterPoolMixin(ExecutorMixin):
+    executor_type = futures.InterpreterPoolExecutor
+
+
 class ProcessPoolForkMixin(ExecutorMixin):
     executor_type = futures.ProcessPoolExecutor
     ctx = "fork"